swf-monitor/scripts/backfill-evgen-events.py

0001 #!/usr/bin/env python3
0002 """backfill-evgen-events.py — event counts for EVGEN datasets registered
0003 before counting existed (docs/RUCIO_REGISTRATION_CONTRACT.md).
0004
0005 Reads the recorded EVGEN inventory snapshot (evgen-rucio.json), selects
0006 the datasets whose Rucio ``events`` is unset, and runs the registration
0007 doer's ``--events-only`` mode on each (register-evgen-rucio.py): the
0008 per-file counts are read through the door and written as the files'
0009 ``events`` attribute, and the dataset's derived total is verified.
0010 Dry run by default (counts, writes nothing); ``--apply`` writes.
0011
0012 Runs under the production operations agent's environment
0013 (EVGEN_X509_PROXY and the door settings); ``--env-file`` loads a
0014 systemd-style KEY=VALUE file when run by hand. After a run, the EVGEN
0015 assimilation ("Update EVGEN from Rucio", or the nightly sweep) carries
0016 the counts into the inventory.
0017
0018 Usage:
0019     backfill-evgen-events.py [--apply] [--env-file /opt/swf-monitor/config/env/production.env]
0020                              [--limit N] [--snapshot PATH]
0021 """
0022 import argparse
0023 import json
0024 import os
0025 import subprocess
0026 import sys
0027 import time
0028
0029 DOER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'register-evgen-rucio.py')
0030 DEFAULT_SNAPSHOT = os.path.join(os.environ.get('SWF_TMP_DIR', '/data/swf-tmp'),
0031                                 'rucio-snapshots', 'evgen-rucio.json')
0032
0033
0034 def load_env_file(path):
0035     for line in open(path):
0036         line = line.strip()
0037         if line and not line.startswith('#') and '=' in line:
0038             key, value = line.split('=', 1)
0039             os.environ.setdefault(key, value.strip().strip('"'))
0040
0041
0042 def main(argv):
0043     ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
0044     ap.add_argument('--apply', action='store_true', help='write the counts to Rucio')
0045     ap.add_argument('--env-file', help='KEY=VALUE file with the agent environment')
0046     ap.add_argument('--snapshot', default=DEFAULT_SNAPSHOT)
0047     ap.add_argument('--limit', type=int, default=0, help='stop after N datasets')
0048     args = ap.parse_args(argv[1:])
0049     if args.env_file:
0050         load_env_file(args.env_file)
0051
0052     with open(args.snapshot) as f:
0053         snap = json.load(f)
0054     todo = [d['did'] for d in snap.get('datasets') or []
0055             if d.get('did') and d.get('events') in (None, 0)]
0056     if args.limit:
0057         todo = todo[:args.limit]
0058     print(f'{len(todo)} of {len(snap.get("datasets") or [])} EVGEN datasets without '
0059           f'a count; {"writing" if args.apply else "dry run"}', flush=True)
0060
0061     results = []
0062     t_all = time.monotonic()
0063     for i, did in enumerate(todo, 1):
0064         path = '/' + did.partition(':')[2].lstrip('/')
0065         cmd = [sys.executable, DOER, '--path', path, '--events-only']
0066         if not args.apply:
0067             cmd.append('--dry-run')
0068         t0 = time.monotonic()
0069         try:
0070             p = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
0071         except subprocess.TimeoutExpired:
0072             results.append({'did': did, 'ok': False, 'error': 'timed out after 3600s'})
0073             print(f'[{i}/{len(todo)}] {did}: TIMEOUT', flush=True)
0074             continue
0075         try:
0076             summary = json.loads((p.stdout or '').strip().splitlines()[-1])
0077         except Exception:                                       # noqa: BLE001
0078             summary = {'ok': False, 'error': (p.stderr or '').strip().splitlines()[-1:]
0079                        or f'rc {p.returncode}'}
0080         rec = {'did': did, 'ok': bool(summary.get('ok')), 'events': summary.get('events'),
0081                'files': summary.get('files'), 'seconds': round(time.monotonic() - t0, 1),
0082                'error': summary.get('error', '')}
0083         rec['failed'] = [x for d in (summary.get('datasets') or []) for x in d.get('failed') or []]
0084         results.append(rec)
0085         print(f'[{i}/{len(todo)}] {did}: {"ok" if rec["ok"] else "FAILED"} '
0086               f'events={rec["events"]} files={rec["files"]} {rec["seconds"]}s'
0087               + (f' {rec["error"]}' if rec['error'] else ''), flush=True)
0088         for x in rec['failed'][:5]:
0089             print(f'      {x["file"]}: {x["reason"]}', flush=True)
0090
0091     ok = sum(1 for r in results if r['ok'])
0092     total_events = sum(r['events'] or 0 for r in results if r['ok'])
0093     print(json.dumps({'datasets': len(results), 'ok': ok, 'failed': len(results) - ok,
0094                       'events_total': total_events, 'apply': args.apply,
0095                       'seconds': round(time.monotonic() - t_all, 1),
0096                       'failures': [{'did': r['did'], 'error': r['error'],
0097                                     'files': r['failed'][:5]}
0098                                    for r in results if not r['ok']]}), flush=True)
0099     return 0 if ok == len(results) else 1
0100
0101
0102 if __name__ == '__main__':
0103     sys.exit(main(sys.argv))