scripts/es-probe/es_range_client.py

0001 #!/usr/bin/env python3
0002 """Event Service probe payload: speak the pilot's range channel.
0003
0004 Runs inside an Event Service job under the pilot's generic executor.
0005 The pilot opens a yampl server socket and exports its name as
0006 PILOT_EVENTRANGECHANNEL; this client connects, asks for ranges, does a
0007 moment of work per range, writes a receipt file per range and reports
0008 each one finished in the pilot's message form:
0009
0010     payload -> pilot   "Ready for events"
0011     pilot -> payload   JSON list of ranges, or "No more events"
0012     payload -> pilot   "<output path>,ID:<eventRangeID>,CPU:<s>,WALL:<s>"
0013
0014 The pilot tars the reported outputs, stages them out on the queue's
0015 es_events/pw activity and reports the ranges finished to the server.
0016 Exit 0 when the pilot says there are no more events.
0017
0018 Usage in the job's exec: python3 es_range_client.py $PILOT_EVENTRANGECHANNEL
0019 (the token on the command line is what makes the pilot export it).
0020
0021 Needs the python-yampl module on PYTHONPATH (built on the host by
0022 swf-epicprod tools/npps0/build-yampl.sh; the pass script exports it and
0023 the pilot's environment reaches the payload).
0024 """
0025
0026 import json
0027 import os
0028 import sys
0029 import time
0030
0031
0032 def log(msg):
0033     print(f"[es_range_client] {time.strftime('%H:%M:%S')} {msg}", flush=True)
0034
0035
0036 def main():
0037     # The pilot hands over the channel by string match: only a payload
0038     # command that names PILOT_EVENTRANGECHANNEL gets the export prefix
0039     # (pilot3 esprocess.py); any other command is given the AthenaMP
0040     # --preExec form. So the job's exec passes $PILOT_EVENTRANGECHANNEL
0041     # as the argument, and the environment is the fallback.
0042     name = sys.argv[1] if len(sys.argv) > 1 and sys.argv[1] else os.environ.get('PILOT_EVENTRANGECHANNEL')
0043     if not name:
0044         log("ERROR: no channel name (argument or PILOT_EVENTRANGECHANNEL); not an Event Service job")
0045         return 2
0046     try:
0047         import yampl
0048     except Exception as exc:  # the host must provide the binding
0049         log(f"ERROR: cannot import yampl: {exc}; PYTHONPATH={os.environ.get('PYTHONPATH')}")
0050         return 3
0051
0052     work_seconds = float(os.environ.get('ES_PROBE_WORK_SECONDS', '2'))
0053     # Receipts go to the job directory, one level above runGen's workDir
0054     # (the current directory): runGen removes workDir when the payload
0055     # exits, before the pilot tars and stages the reported outputs.
0056     outdir = os.path.abspath(os.environ.get('ES_PROBE_OUTDIR',
0057                                             os.path.join(os.path.dirname(os.getcwd()), 'es_probe_out')))
0058     os.makedirs(outdir, exist_ok=True)
0059
0060     sock = yampl.ClientSocket(name, 'local')
0061     log(f"connected to channel {name}; work {work_seconds}s per range; receipts in {outdir}")
0062
0063     done = 0
0064     while True:
0065         sock.send_raw(b"Ready for events")
0066         while True:
0067             size, buf = sock.try_recv_raw()
0068             if size != -1:
0069                 break
0070             time.sleep(0.05)
0071         message = buf.decode('utf8') if isinstance(buf, bytes) else str(buf)
0072         if "No more events" in message:
0073             log(f"no more events after {done} ranges")
0074             return 0
0075         try:
0076             ranges = json.loads(message)
0077         except Exception as exc:
0078             log(f"ERROR: unparseable range message: {exc}: {message[:300]}")
0079             return 4
0080         if isinstance(ranges, dict):
0081             ranges = [ranges]
0082         for rng in ranges:
0083             rid = rng.get('eventRangeID')
0084             t0 = time.time()
0085             c0 = time.process_time()
0086             time.sleep(work_seconds)
0087             wall = time.time() - t0
0088             cpu = time.process_time() - c0
0089             path = os.path.join(outdir, f"{rid}.json")
0090             with open(path, 'w') as fh:
0091                 json.dump({'range': rng, 'wall': wall, 'cpu': cpu,
0092                            'host': os.uname().nodename, 'pid': os.getpid()}, fh)
0093             report = f"{path},ID:{rid},CPU:{cpu:.3f},WALL:{wall:.3f}"
0094             sock.send_raw(report.encode('utf8'))
0095             done += 1
0096             log(f"range {rid} events {rng.get('startEvent')}-{rng.get('lastEvent')} of {rng.get('LFN')} finished: {report}")
0097
0098
0099 if __name__ == '__main__':
0100     sys.exit(main())