Merge branch 'skeletonCA' into 'main'

Minimal working version of CA skeleton for evgen

See merge request !68732

Generators/EvgenJobTransforms/python/EvgenCAConfig.py

+#  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
+
+# Get logger
+from AthenaCommon.Logging import logging
+evgenLog = logging.getLogger('EvgenConfig')
+
+class EvgenConfig():
+    """The CA-based EvgenConfig class that holds the configuration for a sample to be generated"""
+
+    __slots__ = ()
+ 
+    def __init__(self, flags):
+        self.generators = []
+        self.keywords = []
+        self.contact = []
+        self.nEventsPerJob = None
+
+    def setupFlags(self, flags):
+        raise RuntimeError("setupFlags method needs to be implemented in Sample(EvgenConfig)")
+    
+    def checkAttributes(self):
+        for var, value in vars(self).items():
+           if not value:
+               raise RuntimeError("self.{} should be set in Sample(EvgenConfig)".format(var))
+           else:
+               evgenLog.info("{} = {}".format(var, value))
+    
+    def setupProcess(self, flags):
+        raise RuntimeError("setupProcess method needs to be implemented in Sample(EvgenConfig)")

Generators/EvgenJobTransforms/python/GENtoEVGEN_Skeleton.py

@@ -2,8 +2,6 @@
 #
 """Functionality core of the Gen_tf transform"""
 
-from PyJobTransforms.CommonRunArgsToFlags import commonRunArgsToFlags
-
 # temporarily force no global config flags
 from AthenaConfiguration import AllConfigFlags
 del AllConfigFlags.ConfigFlags
@@ -12,32 +10,180 @@ del AllConfigFlags.ConfigFlags
 from AthenaCommon import JobProperties
 JobProperties.jobPropertiesDisallowed = True
 
+# Get logger
+from AthenaCommon.Logging import logging
+evgenLog = logging.getLogger('Gen_tf')
+    
+# Functions for pre/post-include/exec
+from PyJobTransforms.TransformUtils import processPreExec, processPreInclude
+
+# For creating CA instances of algorithms
+from AthenaConfiguration.ComponentFactory import CompFactory
+
+# Other imports that are needed
+import sys, os
+
+
+# Function that reads the jO and returns an instance of Sample(EvgenCAConfig)
+def setupSample(runArgs, flags):
+    # Only permit one jobConfig argument for evgen
+    if len(runArgs.jobConfig) != 1:
+        evgenLog.info("runArgs.jobConfig = %s", runArgs.jobConfig)
+        evgenLog.error("You must supply one and only one jobConfig file argument")
+        sys.exit(1)
+    
+    evgenLog.info("Using JOBOPTSEARCHPATH (as seen in skeleton) = {}".format(os.environ["JOBOPTSEARCHPATH"]))
+    
+    FIRST_DIR = (os.environ['JOBOPTSEARCHPATH']).split(":")[0]
+    
+    # Find jO file
+    jofiles = [f for f in os.listdir(FIRST_DIR) if (f.startswith('mc') and f.endswith('.py'))]
+    if len(jofiles) !=1:
+        evgenLog.error("You must supply one and only one jobOption file in DSID directory")
+        sys.exit(1)
+    jofile = jofiles[0]
+    
+    # Perform consistency checks on the jO
+    from GeneratorConfig.GenConfigHelpers import checkJOConsistency
+    checkJOConsistency(jofile)
+    
+    # Import the jO as a module
+    # We cannot do import BLAH directly since 
+    # 1. the filenames are not python compatible (mc.GEN_blah.py)
+    # 2. the filenames are different for every jO       
+    import importlib.util 
+    spec = importlib.util.spec_from_file_location(
+        name="sample",
+        location=os.path.join(FIRST_DIR,jofile),
+    )
+    jo = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(jo)
+    
+    # Create instance of Sample(EvgenCAConfig)
+    sample = jo.Sample(flags)
+    
+    # Set up the sample properties
+    sample.setupFlags(flags)
+            
+    # Set the random number seed
+    # Need to use logic in EvgenJobTransforms.Generate_dsid_ranseed 
+    
+    # Get DSID
+    dsid = os.path.basename(runArgs.jobConfig[0])
+        
+    # Update the global flags
+    if dsid.isdigit(): 
+        flags.Generator.DSID
+    flags.Input.Files = []
+    flags.Input.RunNumbers = [flags.Generator.DSID]
+    flags.Input.TimeStamps = [0]
+    flags.Generator.nEventsPerJob = sample.nEventsPerJob
+    flags.Exec.MaxEvents = runArgs.maxEvents if runArgs.maxEvents != -1 else sample.nEventsPerJob
+    flags.Output.EVNTFileName = runArgs.outputEVNTFile
+
+    # Check if sample attributes have been properly set
+    sample.checkAttributes()
+
+    return sample
+
+
+# Function to set up event generation services
+def setupEvgenServices(flags, cfg):   
+    
+    # EventInfoCnvAlg
+    from xAODEventInfoCnv.xAODEventInfoCnvConfig import EventInfoCnvAlgCfg
+    cfg.merge(EventInfoCnvAlgCfg(flags, disableBeamSpot=True, xAODKey="TMPEvtInfo"))
+
+
+# Function to set up event generation services
+def setupOutputStream(flags, runArgs, cfg, sample):   
+    
+    # Items to save
+    ItemList = ["McEventCollection#*"]
+    
+    # Configure output stream
+    from OutputStreamAthenaPool.OutputStreamConfig import OutputStreamCfg
+    cfg.merge(
+        OutputStreamCfg(
+            flags, 
+            "EVNT", 
+            ItemList = ItemList
+        )
+    )
+        
+    # Add in-file MetaData
+    from xAODMetaDataCnv.InfileMetaDataConfig import SetupMetaDataForStreamCfg
+    cfg.merge(SetupMetaDataForStreamCfg(flags, "EVNT"))  
+    
+    # Write AMI tag into in-file MetaData
+    from PyUtils.AMITagHelperConfig import AMITagCfg
+    cfg.merge(AMITagCfg(flags, runArgs))    
+
 
 # Main function
 def fromRunArgs(runArgs):
-    from AthenaCommon.Logging import logging
-    evgenLog = logging.getLogger('Gen_tf')
-    evgenLog.setLevel(logging.DEBUG if runArgs.VERBOSE else logging.INFO)    
+    # Set logging level
+    evgenLog.setLevel(logging.DEBUG if runArgs.VERBOSE else logging.INFO)
     
-    # Set up sequences
-        
     # Announce arg checking
-    evgenLog.debug("****************** CHECKING EVENT GENERATION ARGS *****************")
-    evgenLog.debug(runArgs)
-
-    evgenLog.debug('****************** Setting-up configuration flags *****************')
+    if runArgs.VERBOSE:
+        evgenLog.debug("****************** CHECKING EVENT GENERATION ARGS *****************")
+        evgenLog.debug(runArgs)
+        evgenLog.debug('****************** Setting-up configuration flags *****************')
+    
     from AthenaConfiguration.AllConfigFlags import initConfigFlags
     flags = initConfigFlags()
     
-    # Convert run arguments to athena flags
-    from GeneratorConfig.GeneratorConfigFlags import  generatorRunArgsToFlags
+    # Convert run arguments to global athena flags
+    from PyJobTransforms.CommonRunArgsToFlags import commonRunArgsToFlags
     commonRunArgsToFlags(runArgs, flags)
+    
+    # Convert generator-specific run arguments to global athena flags
+    from GeneratorConfig.GeneratorConfigFlags import  generatorRunArgsToFlags
     generatorRunArgsToFlags(runArgs, flags)
+        
+    # convert arguments to flags 
+    flags.fillFromArgs()
     
+    # Create an instance of the Sample(EvgenCAConfig) and update global flags accordingly
+    sample = setupSample(runArgs, flags)
+        
+    # Process pre-include
+    processPreInclude(runArgs, flags)
+
+    # Process pre-exec
+    processPreExec(runArgs, flags)
+
     # Lock flags
     flags.lock()
+    
+    if runArgs.VERBOSE:
+        evgenLog.debug('****************** Dump configuration flags *****************')
+        evgenLog.debug(flags.dump())
         
-    # Announce start of job configuration
-    evgenLog.debug("****************** CONFIGURING EVENT GENERATION *****************")
+        # Announce start of job configuration
+        evgenLog.debug("****************** CONFIGURING EVENT GENERATION *****************")
     
-    evgenLog.info("Hello from GENtoEVGEN_Skeleton")
+    # Main object
+    from AthenaConfiguration.MainServicesConfig import MainEvgenServicesCfg
+    cfg = MainEvgenServicesCfg(flags)  
+    
+    # Set up evgen job with all sequences
+    setupEvgenServices(flags, cfg)
+    
+    # Set up the process
+    AthSequencer = CompFactory.AthSequencer
+    evgenGenSeq = AthSequencer('EvgenGenSeq')
+    cfg.addSequence(evgenGenSeq, parentName='AthAllAlgSeq')
+    cfg.merge(sample.setupProcess(flags), sequenceName="EvgenGenSeq")
+    
+    # Set up the output stream
+    setupOutputStream(flags, runArgs, cfg, sample)
+    
+    # Print ComponentAccumulator components
+    if runArgs.VERBOSE:
+        evgenLog.debug("****************** DUMPING CA CONFIG *****************")
+        evgenLog.debug(cfg.printConfig())
+
+    # Run final ComponentAccumulator
+    sys.exit(not cfg.run().isSuccess())

Generators/GeneratorConfig/python/GenConfigHelpers.py

 #  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
 
+# Get logger
+from AthenaCommon.Logging import logging
+evgenLog = logging.getLogger('GenConfigHelpers')
+
 # Generators providing input events via the LHEF format 
 # (used to determine the input file dummy-naming strategy for C++ generators)
 LHEFGenerators = ["Lhef", # generic name: prefer to use the names below
@@ -93,3 +97,35 @@ def gen_sortkey(genname):
 
     # Return a tuple
     return (genstage,  genname)
+
+# Function to perform consistency check on jO
+def checkJOConsistency(jofile):
+    import os, sys, string
+    
+    joparts = (os.path.basename(jofile)).split(".")
+    # Perform some consistency checks
+    if joparts[0].startswith("mc") and all(c in string.digits for c in joparts[0][2:]):
+        # Check that there are exactly 4 name parts separated by '.': MCxx, DSID, physicsShort, .py
+        if len(joparts) != 3:
+            evgenLog.error(jofile + " name format is wrong: must be of the form mc.<physicsShort>.py: please rename.")
+            sys.exit(1)
+        # Check the length limit on the physicsShort portion of the filename
+        jo_physshortpart = joparts[1]
+        if len(jo_physshortpart) > 50:
+            evgenLog.error(jofile + " contains a physicsShort field of more than 60 characters: please rename.")
+            sys.exit(1)
+        # There must be at least 2 physicsShort sub-parts separated by '_': gens, (tune)+PDF, and process
+        jo_physshortparts = jo_physshortpart.split("_")
+        if len(jo_physshortparts) < 2:
+            evgenLog.error(jofile + " has too few physicsShort fields separated by '_': should contain <generators>(_<tune+PDF_if_available>)_<process>. Please rename.")
+            sys.exit(1)
+        
+        # NOTE: a further check on physicsShort consistency is done below, after fragment loading
+        check_jofiles="/cvmfs/atlas.cern.ch/repo/sw/Generators/MC16JobOptions/scripts"
+        sys.path.append(check_jofiles)
+        from check_jo_consistency import check_naming       
+        if os.path.exists(check_jofiles):
+            check_naming(os.path.basename(jofile))
+        else:
+            evgenLog.error("check_jo_consistency.py not found")
+            sys.exit(1)

Generators/GeneratorConfig/python/GeneratorConfigFlags.py

@@ -19,8 +19,11 @@ def createGeneratorConfigFlags():
     # Events per job
     gencf.addFlag("Generator.nEventsPerJob", 10000)
 
+    # Events per job
+    gencf.addFlag("Generator.DSID", 999999)
+
     # First event
-    gencf.addFlag("Generator.firstEvent", 1)
+    gencf.addFlag("Generator.firstEvent", -1)
     
     # Number of HepMC events to print
     gencf.addFlag("Generator.printEvts", 0)
@@ -43,7 +46,7 @@ def generatorRunArgsToFlags(runArgs, flags):
         flags.Beam.Energy = runArgs.ecmEnergy/2	* GeV 
     else:
         raise RuntimeError("No center of mass energy provided.") 
-    
+        
     if hasattr(runArgs, "ignoreBlackList"):
         flags.Generator.ignoreBlackList = runArgs.ignoreBlackList