swf-monitor/scripts/log-grant.py

0001 #!/usr/bin/env python3
0002 """log-grant.py — a presigned upload grant for one production job's log.
0003
0004 A pilot whose log transfer failed holds the log in the devcloud stage-out
0005 bucket until production operations move it to its dataset (swf-epicprod
0006 docs/LOG_STAGEOUT_FALLBACK.md). The pilot carries no credential: it asks the
0007 grant endpoint, which has checked the job against the PanDA record, and the
0008 production operations agent runs this doer to sign a POST for exactly the
0009 object the log goes to. The grant is written to the shared cache the endpoint
0010 serves it from.
0011
0012 The object key is ``logs/<site>/<log dataset>/<lfn>``. The grant allows that
0013 one key, a size between 1 byte and LOG_MAX_BYTES, and lasts 12 hours.
0014
0015 Standalone; also usable by hand::
0016
0017     ../swf-testbed/.venv/bin/python scripts/log-grant.py --pandaid 3674202 \\
0018         --lfn group.EIC....log.tgz --site UM_GREX_PanDA_1 \\
0019         --dataset group.EIC...._log.121307.121307
0020
0021 The last stdout line is a JSON summary; progress goes to stderr.
0022 Exit codes: 0 ok · 2 bad arguments · 5 no usable credential · 1 other failure.
0023 """
0024
0025 import argparse
0026 import json
0027 import os
0028 import re
0029 import sys
0030 import tempfile
0031 import time
0032
0033 BUCKET = os.environ.get('EPICPROD_STAGEOUT_BUCKET', 'epic-devcloud-stageout')
0034 REGION = os.environ.get('EPICPROD_STAGEOUT_REGION', 'us-east-1')
0035 PROFILE = os.environ.get('EPICPROD_STAGEOUT_AWS_PROFILE', 'epic-stageout')
0036 GRANT_LIFETIME_S = 12 * 3600
0037 LOG_MAX_BYTES = 1024 ** 3
0038 GRANT_DIR = os.path.join(os.environ.get('SWF_TMP_DIR', '/data/swf-tmp'), 'log-grants')
0039
0040 NAME = re.compile(r'^[A-Za-z0-9_.\-]{1,250}$')
0041
0042
0043 def grant_path(pandaid):
0044     return os.path.join(GRANT_DIR, f'{int(pandaid)}.json')
0045
0046
0047 def write_atomic(path, record):
0048     os.makedirs(os.path.dirname(path), mode=0o755, exist_ok=True)
0049     fd, tmp = tempfile.mkstemp(dir=os.path.dirname(path), prefix='.grant-')
0050     try:
0051         with os.fdopen(fd, 'w') as fp:
0052             json.dump(record, fp)
0053         os.chmod(tmp, 0o644)
0054         os.replace(tmp, path)
0055     except BaseException:
0056         if os.path.exists(tmp):
0057             os.unlink(tmp)
0058         raise
0059
0060
0061 def main():
0062     parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
0063     parser.add_argument('--pandaid', type=int, required=True)
0064     parser.add_argument('--lfn', required=True)
0065     parser.add_argument('--site', required=True)
0066     parser.add_argument('--dataset', required=True)
0067     args = parser.parse_args()
0068
0069     for field in ('lfn', 'site', 'dataset'):
0070         if not NAME.match(getattr(args, field)):
0071             print(f'log-grant: {field} {getattr(args, field)!r} is not a plain name', file=sys.stderr)
0072             print(json.dumps({'ok': False, 'error': f'bad {field}'}))
0073             return 2
0074
0075     key = f'logs/{args.site}/{args.dataset}/{args.lfn}'
0076     try:
0077         import boto3
0078         from botocore.exceptions import BotoCoreError, ClientError
0079     except ImportError as exc:
0080         print(f'log-grant: boto3 unavailable: {exc}', file=sys.stderr)
0081         print(json.dumps({'ok': False, 'error': 'boto3 unavailable'}))
0082         return 5
0083     try:
0084         client = boto3.Session(profile_name=PROFILE).client('s3', region_name=REGION)
0085         post = client.generate_presigned_post(
0086             BUCKET, key, Conditions=[['content-length-range', 1, LOG_MAX_BYTES]],
0087             ExpiresIn=GRANT_LIFETIME_S)
0088     except (BotoCoreError, ClientError) as exc:
0089         print(f'log-grant: signing failed with profile {PROFILE}: {exc}', file=sys.stderr)
0090         print(json.dumps({'ok': False, 'error': f'signing failed: {exc}'}))
0091         return 5
0092
0093     record = {'pandaid': args.pandaid, 'lfn': args.lfn, 'key': key,
0094               'url': post['url'], 'fields': post['fields'],
0095               'expires_at': int(time.time()) + GRANT_LIFETIME_S}
0096     write_atomic(grant_path(args.pandaid), record)
0097     print(f'log-grant: {args.pandaid} {key}', file=sys.stderr)
0098     print(json.dumps({'ok': True, 'pandaid': args.pandaid, 'key': key}))
0099     return 0
0100
0101
0102 if __name__ == '__main__':
0103     sys.exit(main())