7from pathlib
import Path
9from AnaAlgorithm.DualUseConfig
import isAthena
10from AnaAlgorithm.Logging
import logging
12logCPGridRun = logging.getLogger(
'CPGridRun')
38 from AnalysisAlgorithmsConfig.AthenaCPRunScript
import AthenaCPRunScript
41 from AnalysisAlgorithmsConfig.EventLoopCPRunScript
import EventLoopCPRunScript
46 parser = argparse.ArgumentParser(description=
'CPGrid runscript to submit CPRun.py jobs to the grid. '
47 'This script will submit a job to the grid using files in the input text one by one.'
48 'CPRun.py can handle multiple sources of input and create one output; but not this script',
50 formatter_class=argparse.RawTextHelpFormatter)
51 parser.add_argument(
'-h',
'--help', dest=
'help', action=
'store_true', help=
'Show this help message and continue')
53 ioGroup = parser.add_argument_group(
'Input/Output file configuration')
54 ioGroup.add_argument(
'-i',
'--input-list', dest=
'input_list', help=
'Path to the text file containing list of containers on the panda grid. Each container will be passed to prun as --inDS and is run individually')
55 ioGroup.add_argument(
'--output-files', dest=
'output_files', nargs=
'+', default=[
'output.root'],
56 help=
'The output files of the grid job. Example: --output-files A.root B.txt B.root results in A/A.root, B/B.txt, B/B.root in the output directory. No need to specify if using CPRun.py')
57 ioGroup.add_argument(
'--destSE', dest=
'destSE', default=
'', type=str, help=
'Destination storage element (PanDA)')
58 ioGroup.add_argument(
'--mergeType', dest=
'mergeType', default=
'Default', type=str, help=
'Output merging type, [None, Default, xAOD]')
60 pandaGroup = parser.add_argument_group(
'Input/Output naming configuration')
61 pandaGroup.add_argument(
'--gridUsername', dest=
'gridUsername', default=os.getenv(
'USER',
''), type=str, help=
'Grid username, or the groupname. Default is the current user. Only affect file naming')
62 pandaGroup.add_argument(
'--prefix', dest=
'prefix', default=
'', type=str, help=
'Prefix for the output directory. Dynamically set with input container if not provided')
63 pandaGroup.add_argument(
'--suffix', dest=
'suffix', default=
'',type=str, help=
'Suffix for the output directory')
64 pandaGroup.add_argument(
'--outDS', dest=
'outDS', default=
'', type=str,
65 help=
'Name of an output dataset. outDS will contain all output files (PanDA). If not provided, support dynamic naming if input name is in the Atlas production format or typical user production format')
67 cpgridGroup = parser.add_argument_group(
'CPGrid configuration')
68 cpgridGroup.add_argument(
'--groupProduction', dest=
'groupProduction', action=
'store_true', help=
'Only use for official production')
70 cpgridGroup.add_argument(
'--exec', dest=
'exec', type=str,
71 help=
'Executable line for the CPRun.py or custom script to run on the grid encapsulated in a double quote (PanDA)\n'
72 'Run CPRun.py with preset behavior including streamlined file i/o. E.g, "CPRun.py -t config.yaml --no-systematics".\n'
73 'Run custom script: "customRun.py -i inputs -o output --text-config config.yaml --flagA --flagB"\n'
76 submissionGroup = parser.add_argument_group(
'Submission configuration')
77 submissionGroup.add_argument(
'--noSubmit', dest=
'noSubmit', action=
'store_true', help=
'Do not submit the job to the grid (PanDA). Useful to inspect the prun command')
78 submissionGroup.add_argument(
'--testRun', dest=
'testRun', action=
'store_true', help=
'Will submit job to the grid but greatly limit the number of files per job (10) and number of events (300)')
79 submissionGroup.add_argument(
'--recreateTar', dest=
'recreateTar', action=
'store_true', help=
'Re-compress the source code. Source code are compressed by default in submission, this is useful when the source code is updated')
80 submissionGroup.add_argument(
'--useCentralPackage', dest=
'useCentralPackage', action=
'store_true', help=
'Use central package instead of custom packages')
81 submissionGroup.add_argument(
'--bulk-submission', dest=
'bulk_submission', action=
'store_true', help=
'Submit all containers in the input list as one task.')
83 miscGroup = parser.add_argument_group(
'Miscellaneous configuration')
84 miscGroup.add_argument(
'-y',
'--agreeAll', dest=
'agreeAll', action=
'store_true', help=
'Agree to all the submission details without asking for confirmation. Use with caution!')
85 miscGroup.add_argument(
'--checkInputDS', dest=
'checkInputDS', action=
'store_true', help=
'Check if the input datasets are available on the AMI.')
86 miscGroup.add_argument(
'--framework', dest=
'framework', default=
'CPGridRun', type=str, help=
'Declaring a name for your submission for PanDA team to collect statistics. Default is CPGridRun')
94 converting unknown args to a dictionary
97 if unknownArgsDict
and self.
hasPrun():
99 logCPGridRun.info(f
"Adding prun exclusive arguments: {unknownArgsDict.keys()}")
100 elif unknownArgsDict:
101 logCPGridRun.warning(f
"Unknown arguments detected: {unknownArgsDict}. Cannot check the availablility in Prun because Prun is not available / noSubmit is on.")
104 return unknownArgsDict
108 if not self.
args.input_list:
109 raise ValueError(
'No input list provided, use --input-list to specify the input containers')
111 input_list_path = Path(self.
args.input_list)
112 if input_list_path.exists():
113 if input_list_path.suffix ==
'.txt':
115 elif input_list_path.suffix ==
'.json':
116 raise NotImplementedError(
'JSON input list parsing is not implemented')
118 raise ValueError(
'Unsupported input list format, only .txt files are supported.')
119 elif CPGridRun.isAtlasProductionFormat(self.
args.input_list):
123 raise ValueError(
'use --input-list to specify input containers')
128 for output
in self.
args.output_files:
130 output_files.extend(output.split(
','))
132 output_files.append(output)
137 logCPGridRun.info(
"\033[92m\n If you are using CPRun.py, the following flags are for the CPRun.py in this framework\033[0m")
138 self.
_runscript.parser.usage = argparse.SUPPRESS
149 self.
cmd[input] = cmd
150 self.
outputs[input] = config[
"outDS"]
157 'cmtConfig': os.environ[
"CMTCONFIG"],
158 'writeInputToTxt':
'IN:in.txt',
162 'addNthFieldOfInDSToLFN':
'2,3,6',
163 'framework': self.
args.framework,
165 if self.
args.noSubmit:
166 config[
'noSubmit'] =
True
168 if self.
args.mergeType ==
'xAOD':
169 config[
'mergeScript'] =
'xAODMerge %OUT `echo %IN | sed \'s/,/ /g\'`'
171 if self.
args.mergeType !=
'None':
172 config[
'mergeOutput'] =
True
175 if self.
args.useCentralPackage:
177 config[
'noBuild'] =
True
178 config[
'noCompile'] =
True
179 config[
'athenaTag'] = f
"AnalysisBase,{os.environ['AnalysisBase_VERSION']}"
181 config[
'outTarBall'] = self.
_tarfile
182 config[
'useAthenaPackages'] =
True
186 config[
'useAthenaPackages'] =
True
188 if self.
args.groupProduction:
189 config[
'official'] =
True
190 config[
'voms'] = f
'atlas:/atlas/{self.args.gridUsername}/Role=production'
193 config[
'destSE'] = self.
args.destSE
195 if self.
args.testRun:
196 config[
'nEventsPerFile'] = 100
200 for k, v
in config.items():
201 if isinstance(v, bool)
and v:
203 elif v
is not None and v !=
'':
204 cmd += f
'--{k} {v} \\\n'
205 return cmd.rstrip(
' \\\n'), config
209 Cleans the unknown args by removing leading dashes and ensuring they are in key-value pairs
211 unknown_args_dict = {}
219 unknown_args_dict[self.
unknown_args[idx].lstrip(
'-')] =
True
221 return unknown_args_dict
225 check the arguments against the prun script to ensure they are valid
226 See https://github.com/PanDAWMS/panda-client/blob/master/pandaclient/PrunScript.py
228 import pandaclient.PrunScript
230 original_argv = sys.argv
233 prunArgsDict = pandaclient.PrunScript.main(get_options=
True)
234 sys.argv = original_argv
235 nonPrunOrCPGridArgs = []
237 if arg
not in prunArgsDict:
238 nonPrunOrCPGridArgs.append(arg)
239 if nonPrunOrCPGridArgs:
240 logCPGridRun.error(f
"Unknown arguments detected: {nonPrunOrCPGridArgs}. They do not belong to CPGridRun or Panda.")
241 raise ValueError(f
"Unknown arguments detected: {nonPrunOrCPGridArgs}. They do not belong to CPGridRun or Panda.")
244 for key, cmd
in self.
cmd.items():
245 parsed_name = CPGridRun.atlasProductionNameParser(key)
246 logCPGridRun.info(
"\n"
248 "\n".join([f
" {k.replace('_', ' ').title()}: {v}" for k, v
in parsed_name.items()]))
249 logCPGridRun.info(f
"Command: \n{cmd}")
257 import pyAMI.atlas.api
258 except ModuleNotFoundError:
260 "Cannot import pyAMI, please run the following commands:\n\n"
263 "voms-proxy-init -voms atlas\n"
265 "and make sure you have a valid certificate.")
273 client = pyAMI.client.Client(
'atlas')
274 pyAMI.atlas.api.init()
278 results = pyAMI.atlas.api.list_datasets(client, patterns=queries)
279 except pyAMI.exception.Error:
281 "Cannot query AMI, please run 'voms-proxy-init -voms atlas' and ensure your certificate is valid.")
288 Helper function to prepare a list of queries for the AMI based on the input list.
289 It will replace the _p### with _p% to match the latest ptag.
292 regex = re.compile(
"_p[0-9]+")
295 for datasetName
in self.
cmd:
296 parsed = CPGridRun.atlasProductionNameParser(datasetName)
297 datasetPtag[datasetName] = parsed.get(
'ptag')
298 queries.append(regex.sub(
"_p%", datasetName))
299 return queries, datasetPtag
303 regex = re.compile(
"_p[0-9]+")
304 results = [r[
'ldn']
for r
in results]
308 for datasetName
in self.
cmd:
309 if datasetName
not in results:
310 notFound.append(datasetName)
312 base = regex.sub(
"_p%", datasetName)
313 matching = [r
for r
in results
if r.startswith(base.replace(
"_p%",
""))]
315 mParsed = CPGridRun.atlasProductionNameParser(m)
317 mPtagInt = int(mParsed.get(
'ptag',
'p0')[1:])
318 currentPtagInt = int(datasetPtag.get(datasetName,
'p0')[1:])
319 if mPtagInt > currentPtagInt:
320 latestPtag[datasetName] = f
"p{mPtagInt}"
321 except (ValueError, TypeError):
325 logCPGridRun.info(
"Newer version of datasets found in AMI:")
326 for name, ptag
in latestPtag.items():
327 logCPGridRun.info(f
"{name} -> ptag: {ptag}")
330 logCPGridRun.error(
"Some input datasets are not available in AMI, missing datasets are likely to fail on the grid:")
331 logCPGridRun.error(
", ".join(notFound))
343 if CPGridRun.isAtlasProductionFormat(name)
and not label:
350 {group/user}.{username}.{prefix}.{DSID}.{format}.{tags}.{suffix}
352 nameParser = CPGridRun.atlasProductionNameParser(name)
353 base =
'group' if self.
args.groupProduction
else 'user'
354 username = self.
args.gridUsername
355 dsid = nameParser[
'DSID']
356 tags =
'_'.join(nameParser[
'tags'])
357 fileFormat = nameParser[
'format']
358 base =
'group' if self.
args.groupProduction
else 'user'
359 prefix = self.
args.prefix
if self.
args.prefix
else nameParser[
'main'].
split(
'_')[0]
362 result = [base, username, prefix, dsid, fileFormat, tags, suffix]
363 return ".".join(filter(
None, result))
367 {group/user}.{username}.{prefix}.{main}.{suffix}
369 parts = name.split(
'.')
370 base =
'group' if self.
args.groupProduction
else 'user'
371 username = self.
args.gridUsername
372 main = label
if label
else parts[2]
373 main = f
"{self.args.prefix}.{main}" if self.
args.prefix
else main
374 main = f
"{main}.{self.args.suffix}" if self.
args.suffix
else main
376 result = [base, username, main]
377 return ".".join(filter(
None, result))
381 return self.
args.suffix
382 if self.
args.testRun:
384 return f
"test_{uuid.uuid4().hex[:6]}"
389 tarball_mtime = os.path.getmtime(self.
_tarfile)
if os.path.exists(self.
_tarfile)
else 0
394 for root, _, files
in os.walk(buildDir):
396 file_path = os.path.join(root, file)
398 if os.path.getmtime(file_path) > tarball_mtime:
399 logCPGridRun.info(f
"File {file_path} is newer than the tarball.")
401 except FileNotFoundError:
405 if sourceDir
is None:
406 logCPGridRun.warning(
"Source directory is not detected, auto-compression is not performed. Use --recreateTar to update the submission")
408 for root, _, files
in os.walk(sourceDir):
410 file_path = os.path.join(root, file)
412 if os.path.getmtime(file_path) > tarball_mtime:
413 logCPGridRun.info(f
"File {file_path} is newer than the tarball.")
415 except FileNotFoundError:
420 buildDir = os.environ[
"CMAKE_PREFIX_PATH"]
421 buildDir = os.path.dirname(buildDir.split(
":")[0])
425 cmakeCachePath = os.path.join(self.
_buildDir(),
'CMakeCache.txt')
427 if not os.path.exists(cmakeCachePath):
429 with open(cmakeCachePath,
'r')
as cmakeCache:
430 for line
in cmakeCache:
431 if '_SOURCE_DIR:STATIC=' in line:
432 sourceDir = line.split(
'=')[1].
strip()
437 if not self.
args.exec:
438 raise ValueError(
'No exec command provided, use --exec to specify the command to run on the grid')
441 isCPRunDefault = self.
args.exec.startswith(
'-')
or self.
args.exec.startswith(
'CPRun.py')
443 'input_list':
'in.txt',
444 'merge_output_files': len(self.
args.output_files) == 1,
446 if not isCPRunDefault:
447 if self.
_isFirstRun: logCPGridRun.warning(
"Non-CPRun.py is detected, please ensure the exec string is formatted correctly. Exec string will not be automatically formatted.")
448 return f
'"{self.args.exec}"'
452 runscriptArgs, unknownArgs = self.
_runscript.parser.parse_known_args(self.
args.exec.split(
' '))
455 unknown_flags = [arg
for arg
in unknownArgs
if arg.startswith(
'--')]
457 logCPGridRun.error(f
"Unknown flags detected in the exec string: {unknown_flags}. Please check the exec string.")
458 raise ValueError(f
"Unknown arguments detected: {unknown_flags}")
461 for key, value
in formatingClause.items():
462 if hasattr(runscriptArgs, key):
463 old_value = getattr(runscriptArgs, key)
464 if old_value
is None or old_value == self.
_runscript.parser.get_default(key):
465 setattr(runscriptArgs, key, value)
466 if self.
_isFirstRun: logCPGridRun.info(f
"Setting '{key}' to '{value}' (CPRun.py default is: '{old_value}')")
468 if self.
_isFirstRun: logCPGridRun.warning(f
"Preserving user-defined '{key}': '{old_value}', default formatting '{value}' will not be applied.")
470 logCPGridRun.error(f
"Formatting clause '{key}' is not recognized in the CPRun.py script. Check CPGridRun.py")
471 raise ValueError(f
"Formatting clause '{key}' is not recognized in the CPRun.py script. Check CPGridRun.py")
474 arg_string =
' '.join(
475 f
'--{k.replace("_", "-")}' if isinstance(v, bool)
and v
else
476 f
'--{k.replace("_", "-")} {v}' for k, v
in vars(runscriptArgs).items()
if v
not in [
None,
False]
478 return f
'"CPRun.py {arg_string}"'
481 from AnalysisAlgorithmsConfig.CPBaseRunner
import CPBaseRunner
482 if not hasattr(runscriptArgs,
'text_config'):
483 self.
_errorCollector[
'no yaml'] =
"No YAML configuration file is specified in the exec string. Please provide one using --text-config"
485 yamlPath = getattr(runscriptArgs,
'text_config')
487 haveLocalYaml = CPBaseRunner.findLocalPathYamlConfig(yamlPath)
489 logCPGridRun.warning(
"A path to a local YAML configuration file is found, but it may not be grid-usable.")
491 repoYamls, _ = CPBaseRunner.findRepoPathYamlConfig(yamlPath)
492 if repoYamls
and len(repoYamls) > 1:
493 self.
_errorCollector[
'ambiguous yamls'] = f
'Multiple files named \"{yamlPath}\" found in the analysis repository. Please provide a more specific path to the config file.\nMatches found:\n' +
'\n'.join(repoYamls)
495 elif repoYamls
and len(repoYamls) == 1:
496 logCPGridRun.info(f
"Found a grid-usable YAML configuration file in the analysis repository: {repoYamls[0]}")
499 if haveLocalYaml
and self.
args.useCentralPackage:
500 logCPGridRun.warning(
"A path to a local YAML configuration file is found, no custom packages are found, proceed with /cvmfs packages only.")
502 if not repoYamls
and not self.
args.useCentralPackage:
503 self.
_errorCollector[
'no usable yaml'] = f
"Grid usable YAML configuration file not found: {yamlPath}"
505 self.
_errorCollector[
'have local yaml'] = f
"Only a local YAML configuration file is found: {yamlPath}, not usable in the grid.\n" \
506 f
"Make sure the YAML file is in build/x86_64-el9-gcc14-opt/data/package_name/config.yaml. You can install the YAML file through CMakeList.txt with `atlas_install_data( data/* )`; use `-t package_name/config.yaml` in the --exec\n"\
507 f
"Or if you are only using central packages, please use the `--useCentralPackage` flag."
510 outputs = [f
'{output.split(".")[0]}:{output}' if ":" not in output
else output
for output
in self.
args.output_files]
511 return ','.join(outputs)
515 prun_path = shutil.which(
"prun")
516 if prun_path
is None:
518 "The 'prun' command is not found. If you are on lxplus, please run the following commands:\n\n"
521 "voms-proxy-init -voms atlas\n"
523 "Make sure you have a valid certificate."
530 for key, cmd
in self.
cmd.items():
531 logCPGridRun.info(f
"Submitting: {self.outputs[key]}")
532 process = subprocess.Popen(cmd, shell=
True, stdout=sys.stdout, stderr=sys.stderr)
533 process.communicate()
538 name = name.split(
":")[1]
540 if name.startswith(
"mc")
or name.startswith(
"data"):
543 logCPGridRun.warning(
"Name is not in the Atlas production format, assuming it is a user production")
549 The custom name has many variations, but most of them follow user/group.username.datasetname.suffix
552 parts = filename.split(
'.')
553 result[
'userType'] = parts[0]
554 result[
'username'] = parts[1]
555 result[
'main'] = parts[2]
556 result[
'suffix'] = parts[-1]
562 Parsing file name into a dictionary, an example is given here
563 mc20_13TeV.410470.PhPy8EG_A14_ttbar_hdamp258p75_nonallhad.deriv.DAOD_PHYS.e6337_s3681_r13167_p5855/DAOD_PHYS.34865530._000740.pool.root.1
565 datasetName: mc20_13TeV.410470.PhPy8EG_A14_ttbar_hdamp258p75_nonallhad.deriv.DAOD_PHYS.e6337_s3681_r13167_p5855
566 projectName: mc20_13TeV
570 main: PhPy8EG_A14_ttbar_hdamp258p75_nonallhad
571 TODO generator: PhPy8Eg
572 TODO tune: A14 # For Pythia8
574 TODO hdamp: 258p75 # For Powheg
575 TODO decayType: nonallhad
578 tags: e###_s###_r###_p###_a###_t###_b#
579 etag: e6337 # EVNT (EVGEN) production and merging
580 stag: s3681 # Geant4 simulation to produce HITS and merging!
581 rtag: r13167 # Digitisation and reconstruction, as well as AOD merging
582 ptag: p5855 # Production of NTUP_PILEUP format and merging
583 atag: aXXX: atlfast configuration (both simulation and digit/recon)
584 ttag: tXXX: tag production configuration
585 btag: bXXX: bytestream production configuration
598 datasetPart, filePart = filename.split(
'/')
600 datasetPart = filename
604 if ':' in datasetPart:
605 datasetPart = datasetPart.split(
':')[1]
608 if datasetPart.startswith(
'user')
or datasetPart.startswith(
'group'):
609 result[
'datasetName'] = datasetPart
613 datasetParts = datasetPart.split(
'.')
614 result[
'datasetName'] = datasetPart
616 result[
'projectName'] = datasetParts[0]
618 campaign_energy = result[
'projectName'].
split(
'_')
619 result[
'campaign'] = campaign_energy[0]
620 result[
'energy'] = campaign_energy[1]
623 result[
'DSID'] = datasetParts[1]
624 result[
'main'] = datasetParts[2]
625 result[
'step'] = datasetParts[3]
626 result[
'format'] = datasetParts[4]
629 tags = datasetParts[5].
split(
'_')
630 result[
'tags'] = tags
632 if tag.startswith(
'e'):
634 elif tag.startswith(
's'):
636 elif tag.startswith(
'r'):
638 elif tag.startswith(
'p'):
640 elif tag.startswith(
'a'):
642 elif tag.startswith(
't'):
644 elif tag.startswith(
'b'):
649 fileParts = filePart.split(
'.')
650 result[
'jediTaskID'] = fileParts[1]
651 result[
'fileNumber'] = fileParts[2]
652 result[
'version'] = fileParts[-1]
658 with path.open(
'r')
as inputText:
659 for line
in inputText.readlines():
661 if line.startswith(
"#")
or not line.strip():
663 files += line.split(
",")
665 files = [file.strip()
for file
in files]
669 if any((path.parent / file).
exists()
or (path.parent / f
"{file}.txt").
exists()
for file
in files):
673 file_path = path.parent / file
674 if not file_path.exists():
675 file_path = path.parent / f
"{file}.txt"
676 if not file_path.exists():
677 logCPGridRun.error(f
"File {file} or {file}.txt does not exist in the input list directory.")
678 raise FileNotFoundError(f
"File {file} or {file}.txt does not exist in the input list directory.")
679 files_current, names_current = CPGridRun._parseInputFileList(file_path, bulk_submission=
True)
680 files_bulk.extend(files_current)
681 names_bulk.extend(names_current)
682 return files_bulk, names_bulk
684 return [
','.join(files)], [path.stem.replace(
"+",
"")]
690 logCPGridRun.error(
"Errors were collected during the script execution:")
693 logCPGridRun.error(f
"{key}: {value}")
694 logCPGridRun.error(
"Please fix the errors and try again.")
699 if self.
args.checkInputDS:
703 if self.
args.agreeAll:
704 logCPGridRun.info(
"You have agreed to all the submission details. Jobs will be submitted without confirmation.")
707 answer = input(
"Please confirm ALL the submission details are correct before submitting [y/n]: ")
708 if answer.lower() ==
'y':
710 elif answer.lower() ==
'n':
711 logCPGridRun.info(
"Feel free to report any unexpected behavior to the CPAlgorithms team!")
713 logCPGridRun.error(
"Invalid input. Please enter 'y' or 'n'. Jobs are not submitted.")
715if __name__ ==
'__main__':
717 cpgrid.configureSubmission()
718 cpgrid.printInputDetails()
719 cpgrid.checkExternalTools()
720 cpgrid.printDelayedErrorCollection()
721 cpgrid.askSubmission()
void print(char *figname, TCanvas *c1)
dict _createPrunArgsDict(self)
outputDSFormatter(self, name, label)
rucioCustomNameParser(filename)
_checkYamlExists(self, runscriptArgs)
_parseGridArguments(self)
bool checkInputInPyami(self)
bool _analyzeAmiResults(self, results, datasetPtag)
isAtlasProductionFormat(name)
tuple[list[str], list[str]] _parseInputFileList(Path path, bool bulk_submission=False)
_prepareAmiQueryFromInputList(self)
configureSubmissionSingleSample(self, input, name)
_filesChangedOrTarballNotCreated(self)
tuple[list[str], list[str]] _inputNames
_hasCompressedTarball(self)
atlasProductionNameParser(filename)
_outputDSFormatter(self, name)
_checkPrunArgs(self, argDict)
configureSubmission(self)
printDelayedErrorCollection(self)
_customOutputDSFormatter(self, name, label)
dict _unknownArgsDict(self)
bool exists(const std::string &filename)
does a file exist
std::vector< std::string > split(const std::string &s, const std::string &t=":")