/
githubmirror
/
cmssw
Обзор
Документация
Войти
/
githubmirror
/
cmssw
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
DQMServices/StreamerIO/scripts/personalPlayback.py
539 строк
19 KB
Gabriel Moreira
refactor: update personalPlayback to handle optional extra arguments
29 апр 2026, 17:30
29 апр 2026, 17:30
db80211
Код
Авторство
О чём код?
#!/usr/bin/env python3 import os import sys import argparse import asyncore import pickle import logging import subprocess import shutil import re import collections import json import tempfile import signal import time import glob # Utilities log_format = '%(asctime)s: %(name)-20s - %(levelname)-8s - %(message)s' logging.basicConfig(format=log_format, level=logging.INFO) root_log = logging.getLogger() class Applet(object): def __init__(self, name, opts, **kwargs): self.name = name self.opts = opts self.kwargs = kwargs self.do_init() def write(self, fp): self.control_fp = fp with open(fp, "wb") as f: pickle.dump(self, f) self.log.info("Written control file: %s", fp) @staticmethod def read(fp): with open(fp, "rb") as f: return pickle.load(f) @property def log(self): return logging.getLogger(self.name) def do_init(self): pass def do_exec(self): pass def preexec_kill_on_pdeath(): import ctypes libc = ctypes.CDLL("libc.so.6") PR_SET_PDEATHSIG = 1 libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) # Actual implementation of the workers class Playback(Applet): re_pattern = re.compile(r'run([0-9]+)_ls([0-9]+)_stream([A-Za-z0-9]+)_([A-Za-z0-9_-]+)\.jsn') def discover_files(self): self.lumi_found = {} files_found = set() streams_found = set() run_found = None for f in os.listdir(self.input): r = self.re_pattern.match(f) if r: run, lumi, stream, stream_source = r.groups() run, lumi = int(run), int(lumi) if run_found is None: run_found = run elif run_found != run: raise Exception("Files from multiple runs are not (yet) supported for as playback input.") lumi_dct = self.lumi_found.setdefault(lumi, { 'streams': {} }) lumi_dct["streams"][stream] = (f, stream_source) files_found.add(f) streams_found.add(stream) if run_found is None: raise Exception("Playback files not found.") if self.run < 0: self.run = run_found self.log.info("Found run %s, will map output to run %s", run_found, self.run) self.log.info("Found %d lumisections with %d files", len(self.lumi_found), len(files_found)) self.log.info("Found %d streams: %s", len(streams_found), list(streams_found)) self.lumi_order = list(self.lumi_found.keys()) self.lumi_order.sort() self.log.info("Lumi order: %s", str(self.lumi_order)) def do_init(self): # check if our input directory is okay self.input = self.opts.playback self.ramdisk = self.opts.work_ramdisk self.run = self.opts.run self.log.info("Using input directory: %s", self.input) self.discover_files() self.output = os.path.join(self.ramdisk, "run%06d" % self.run) if not os.path.isdir(self.output): os.makedirs(self.output) self.log.info("Using output directory: %s", self.output) self.global_file = os.path.join(self.ramdisk, ".run%06d.global" % self.run) self.log.info("Writing: %s", self.global_file) with open(self.global_file, "w") as f: f.write("run_key = pp_run") self.lumi_backlog = collections.deque() self.lumi_backlog_size = 10 self.next_lumi_index = 1 def do_create_lumi(self): orig_lumi = self.lumi_order[(self.next_lumi_index - 1) % len(self.lumi_order)] play_lumi = self.next_lumi_index self.next_lumi_index += 1 self.log.info("Start copying lumi (original) %06d -> %06d (playback)", orig_lumi, play_lumi) lumi_dct = self.lumi_found[orig_lumi] streams = lumi_dct["streams"] def ijoin(f): return os.path.join(self.input, f) def ojoin(f): return os.path.join(self.output, f) written_files = set() for stream, v in streams.items(): jsn_orig_fn, stream_source = v jsn_play_fn = "run%06d_ls%04d_stream%s_%s.jsn" % (self.run, play_lumi, stream, stream_source) # define dat filename ext = "dat" if stream.startswith("streamDQMHistograms"): ext = "pb" dat_play_fn = "run%06d_ls%04d_stream%s_%s.%s" % (self.run, play_lumi, stream, stream_source, ext) # read the original file name, for copying with open(ijoin(jsn_orig_fn), 'r') as f: jsn_data = json.load(f) dat_orig_fn = jsn_data["data"][3] # copy the data file if os.path.exists(ijoin(dat_orig_fn)): self.log.info("C: %s -> %s", dat_orig_fn, dat_play_fn) shutil.copyfile(ijoin(dat_orig_fn), ojoin(dat_play_fn)) written_files.add(dat_play_fn) else: self.log.warning("Dat file is missing: %s", dat_orig_fn) # write a new json file point to a different data file # this has to be atomic! jsn_data["data"][3] = dat_play_fn f = tempfile.NamedTemporaryFile(prefix=jsn_play_fn+ ".", suffix=".tmp", dir = self.output, delete=False, mode='w+') tmp_fp = f.name self.log.info(jsn_data) json.dump(jsn_data, f) f.close() os.rename(tmp_fp, ojoin(jsn_play_fn)) written_files.add(jsn_play_fn) self.log.info("Copied %d files for lumi %06d", len(written_files), play_lumi) self.lumi_backlog.append((play_lumi, written_files)) while len(self.lumi_backlog) > self.lumi_backlog_size: old_lumi, files_to_delete = self.lumi_backlog.popleft() self.log.info("Deleting %d files for old lumi %06d", len(files_to_delete), old_lumi) for f in files_to_delete: os.unlink(ojoin(f)) def do_exec(self): last_write = 0 lumi_produced = 0 while True: time.sleep(1) now = time.time() if (now - last_write) > self.opts.playback_time_lumi: last_write = now if self.opts.playback_nlumi > -1 and lumi_produced >= self.opts.playback_nlumi: break self.do_create_lumi() lumi_produced += 1 # write eor eor_fn = "run%06d_ls0000_EoR.jsn" % (self.run, ) eor_fp = os.path.join(self.output, eor_fn) with open(eor_fp, "w"): pass self.log.info("Wrote EoR: %s", eor_fp) start_dqm_job = """ #!/bin/env /bin/bash set -x #echo on TODAY=$(date) logname="/var/log/hltd/pid/hlt_run$4_pid$$.log" lognamez="/var/log/hltd/pid/hlt_run$4_pid$$_gzip.log.gz" #override the noclobber option by using >| operator for redirection - then keep appending to log echo startDqmRun invoked $TODAY with arguments $@ >| $logname export http_proxy="http://cmsproxy.cms:3128" export https_proxy="https://cmsproxy.cms:3128/" export NO_PROXY=".cms" export SCRAM_ARCH=$2 cd $1 cd base source cmsset_default.sh >> $logname cd $1 cd current pwd >> $logname 2>&1 eval `scram runtime -sh`; cd $3; pwd >> $logname 2>&1 #exec esMonitoring.py -z $lognamez cmsRun `readlink $6` runInputDir=$5 runNumber=$4 ${{@:7}} >> $logname 2>&1 exec esMonitoring.py cmsRun `readlink $6` runInputDir=$5 runNumber=$4 ${{@:7}} """ start_dqm_job = start_dqm_job.replace("/var/log/hltd/pid/", '{log_path}/') start_dqm_job = start_dqm_job.replace(" cmsRun ", ' {cmsRun} ') RunDesc = collections.namedtuple('Run', ['run', 'run_fp', 'global_fp', 'global_param']) RunState = collections.namedtuple('RunState', ['desc', 'proc']) class FrameworkJob(Applet): def _set_name(self): x = os.path.basename(self.cfg_file) x = re.sub(r'(.*)\.py', r'\1', x) x = re.sub(r'(.*)_cfg', r'\1', x) x = re.sub(r'(.*)-live', r'\1', x) x = re.sub(r'(.*)_sourceclient', r'\1', x) x = re.sub(r'(.*)_dqm', r'\1', x) x = "".join([c for c in x if c.isalnum()]) self.tag = x self.name = "cmssw_%s" % x def _find_release(self): fp = os.path.realpath(self.cfg_file) while len(fp) > 3: bn = os.path.basename(fp) fp = os.path.dirname(fp) if bn == "src": break if len(fp) <= 3: raise Exception("Could not find the cmssw release area.") self.cmsenv_path = fp self.log.info("cmsenv path: %s", self.cmsenv_path) def _prepare_files(self): self.home_path = os.path.join(self.opts.work_home, "%s_%s" % (self.name, hex(id(self)))) self.home_path = os.path.realpath(self.home_path) os.makedirs(self.home_path) self.log_path = self.opts.work_logs self.log.info("logs path: %s", self.log_path) self.exec_file = os.path.join(self.home_path, "startDqmRun.sh") self.log.info("Creating: %s", self.exec_file) f = open(self.exec_file, "w") template = start_dqm_job body = template.format(log_path=self.log_path, cmsRun=self.opts.cmsRun) f.write(body) f.close() os.chmod(self.exec_file, 0o755) cmsset_globs = ["/cvmfs/cms.cern.ch/cmsset_default.sh", "/afs/cern.ch/cms/cmsset_default.sh", "/home/dqm*local/base/cmsset_default.sh"] cmsset_target = None for t in cmsset_globs: files = glob.glob(t) for f in files: cmsset_target = f break if cmsset_target is not None: base = os.path.join(self.home_path, "base") os.makedirs(base) cmsset_link = os.path.join(base, "cmsset_default.sh") self.log.info("Linking : %s -> %s", cmsset_link, cmsset_target) os.symlink(cmsset_target, cmsset_link) else: self.log.warning("Couldn't find cmsset_default.sh, source it yourself!") current_link = os.path.join(self.home_path, "current") target = os.path.relpath(self.cmsenv_path, self.home_path) self.log.info("Linking : %s -> %s", current_link, target) os.symlink(target, current_link) # check if current is outside the release directory # otherwise scram gets stuck forever cp = os.path.commonprefix([self.home_path, self.cmsenv_path]) if self.cmsenv_path == cp: self.log.error("Working directory (incl. control directory), have to be outside the cmssw release. Otherwise scram fails due to recursive links.") raise Exception("Invalid home_path: %s" % self.home_path) output_link = os.path.join(self.home_path, "output") output_target = os.path.realpath(self.opts.work_output) target = os.path.relpath(output_target, self.home_path) self.log.info("Linking : %s -> %s", output_link, target) os.symlink(target, output_link) os.makedirs(os.path.join(output_link, "upload"), exist_ok=True) self.output_path = output_link cfg_link = os.path.join(self.home_path, os.path.basename(self.cfg_file)) target = self.cfg_fp self.log.info("Linking : %s -> %s", cfg_link, target) os.symlink(target, cfg_link) self.cfg_link = cfg_link def do_init(self): # check if our input directory is okay self.ramdisk = self.opts.work_ramdisk self.run = self.opts.run self.cfg_file = self.kwargs["cfg_file"] if not os.path.isfile(self.cfg_file): raise Exception("Configuration file not found: %s" % self.cfg_file) self.cfg_fp = os.path.realpath(self.cfg_file) self.ramdisk_fp = os.path.realpath(self.ramdisk) self._set_name() self._find_release() self._prepare_files() def make_args(self, run): args = [] args.append("bash") # arg 0 args.append(self.exec_file) # arg 0 args.append(self.home_path) # home path args.append(self.opts.arch) # release args.append(self.output_path) # cwd/output path args.append(str(run)) # run args.append(self.ramdisk_fp) # ramdisk args.append(self.cfg_link) # cmsRun arg 1 args.append(f"runkey={self.opts.run_key}") # cmsRun arg 2 for extra_arg in self.opts.extra_args: args.append(extra_arg) return args def discover_latest(self): re_run = re.compile(r'run([0-9]+)') re_global = re.compile(r'\.run([0-9]+)\.global') # find runs runs = {} globals = {} for x in os.listdir(self.ramdisk): m = re_run.match(x) if m: runs[int(m.group(1))] = x m = re_global.match(x) if m: globals[int(m.group(1))] = x # find max global for which there is a run directory run_set = set(runs.keys()) run_set = run_set.intersection(globals.keys()) if self.opts.run < 0: largest = max(run_set) else: largest = self.opts.run #self.log.info("Largest: %s", largest) global_fp = os.path.join(self.ramdisk, globals[largest]) with open(global_fp, "r") as f: global_param = f.read() return RunDesc( run = largest, run_fp = os.path.join(self.ramdisk, runs[largest]), global_fp = global_fp, global_param = global_param, ) def start_run(self, current): old_state = self.current_state # kill the old run # nope, since it involves eof and i am lazy if old_state: return args = self.make_args(current.run) self.log.info("Executing: %s", " ".join(args)) proc = subprocess.Popen(args, preexec_fn=preexec_kill_on_pdeath) self.current_state = RunState(desc=current, proc=proc) def do_exec(self): time.sleep(1) self.current_state = None while True: latest = self.discover_latest() if self.current_state is None or latest != self.current_state.desc: self.log.info("Found latest run: %s", latest) self.start_run(latest) if not self.current_state: self.log.info("Run not found, waiting 1 sec.") else: r = self.current_state.proc.poll() if r is not None: self.log.info("Process exitted: %s", r) return 0 time.sleep(1) import getpass if __name__ == "__main__": if len(sys.argv) == 2 and sys.argv[-1].endswith(".pkl"): f = sys.argv[-1] obj = Applet.read(f) ret = obj.do_exec() sys.exit(ret if ret else 0) # control -> interal files and home directory for the run subdirectories = ["ramdisk", "output", "control", "home", "logs", "dqm_monitoring"] username = getpass.getuser() parser = argparse.ArgumentParser(description="Emulate DQM@P5 environment and launch cmssw jobs.") #parser.add_argument('-q', action='store_true', help="Don't write to stdout, just the log file.") #parser.add_argument("log", type=str, help="Filename to write.", metavar="<logfile.gz>") parser.add_argument("--work", "-w", type=str, help="Working directory (used for inputs,outputs,monitoring and logs).", default="/tmp/pplay." + username) parser.add_argument("--clean", "-c", action="store_true", help="Clean work directories (if they are not set).", default=False) parser.add_argument("--dry", "-n", action="store_true", help="Do not execute, just init.", default=False) parser.add_argument("--arch", "-a", type=str, help="scram architecture to use", default=os.environ.get("SCRAM_ARCH")) work_group = parser.add_argument_group('Paths', 'Path options for cmssw jobs, auto generated if not specified.') for subdirectory in subdirectories: work_group.add_argument("--work_%s" % subdirectory, type=str, help="Path for %s directory." % subdirectory, default=None) playback_group = parser.add_argument_group('Playback', 'Playback configuration/parameters.') playback_group.add_argument("--playback", "-p", type=str, metavar="PLAYBACK_INPUT_DIR", help="Enable playback (emulate file delivery, otherwise set work_input).", default=None) playback_group.add_argument("--playback_nlumi", type=int, help="Number of lumis to deliver, -1 for forever.", default=-1) playback_group.add_argument("--playback_time_lumi", type=float, help="Number of seconds between lumisections.", default=23.3) run_group = parser.add_argument_group('Run', 'Run configuration/parameters.') run_group.add_argument("--run", type=int, help="Run number, -1 for autodiscovery.", default=-1) run_group.add_argument("--cmsRun", type=str, help="cmsRun command to run, for igprof and so on.", default="cmsRun") run_group.add_argument("--run_key", type=str, help="Run key to use (pp_run, cosmic_run or hi_run)", default="pp_run") client_group = parser.add_argument_group("Client", "Client extra arguments.") client_group.add_argument("--extra-args", nargs="+", metavar="KEY=VALUE", help="Client extra arguments split by whitespace") parser.add_argument('cmssw_configs', metavar='cmssw_cfg.py', type=str, nargs='*', help='List of cmssw jobs (clients).') args = parser.parse_args() if len(args.cmssw_configs) and args.cmssw_configs[0] == "--": # compat with 2.6 args.cmssw_configs = args.cmssw_configs[1:] for subdirectory in subdirectories: if getattr(args, "work_" + subdirectory) is None: setattr(args, "work_" + subdirectory, os.path.join(args.work, subdirectory)) path = getattr(args, "work_" + subdirectory) if args.clean and os.path.isdir(path): root_log.info("Removing directory: %s", path) shutil.rmtree(path) path = getattr(args, "work_" + subdirectory) if not os.path.isdir(path): os.makedirs(path) root_log.info("Using directory: %s", path) print("*"*80) print(args) print("*"*80) applets = [] if args.playback: # launch playback service playback = Playback("playback_emu", opts=args) applets.append(playback) for cfg in args.cmssw_configs: cfg_a = FrameworkJob("framework_job", opts=args, cfg_file=cfg) applets.append(cfg_a) if len(applets) == 0: sys.stderr.write("At least one process should be specified, use --playback and/or cmssw_configs options.\n") # serialize them into control directory for a in applets: fn = "%s_%s.pkl" % (a.name, hex(id(a))) a.write(os.path.join(args.work_control, fn)) if args.dry: sys.exit(0) # launch each in a separate subprocess for a in applets: fp = a.control_fp args = [os.path.realpath(__file__), fp] a.control_proc = subprocess.Popen(args, preexec_fn=preexec_kill_on_pdeath) for a in applets: # wait till everything finishes a.control_proc.wait()