karabo-bridge-serve-run command

docs/aligning_trains.ipynb

@@ -849,7 +849,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "xfel",
+   "display_name": "xfel (Python 3.7)",
    "language": "python",
    "name": "xfel"
   },

extra_data/cli/serve_run.py

@@ -0,0 +1,88 @@
+from argparse import ArgumentParser
+import sys
+
+from .. import open_run
+
+IMPORT_FAILED_MSG = """\
+{}
+
+karabo-bridge-serve-run requires additional dependencies:
+    pip install karabo-bridge psutil
+"""
+
+def main(argv=None):
+    ap = ArgumentParser(prog="karabo-bridge-serve-run")
+    ap.add_argument("proposal", help="Proposal number")
+    ap.add_argument("run", help="Run number")
+    ap.add_argument("port", help="TCP port or ZMQ endpoint to send data on")
+    ap.add_argument(
+        "--include", help="Select matching sources (and optionally keys) to "
+                          "include in streamed data",
+        action='append'
+    )
+    ap.add_argument(
+        "--allow-partial", help="Send trains where some sources are missing",
+        action='store_true'
+    )
+    ap.add_argument(
+        "--append-detector-modules", help="combine multiple module sources"
+        " into one (will only work for AGIPD data currently).",
+        action='store_true'
+    )
+    ap.add_argument(
+        "--dummy-timestamps", help="create dummy timestamps if the meta-data"
+        " lacks proper timestamps",
+        action='store_true'
+    )
+    ap.add_argument(
+        "--use-infiniband", help="Use infiniband interface if available "
+                                 "(if a TCP port is specified)",
+        action='store_true'
+    )
+    ap.add_argument(
+        "-z", "--socket-type", help="ZeroMQ socket type",
+        choices=['PUB', 'PUSH', 'REP'], default='REP'
+    )
+    args = ap.parse_args(argv)
+
+    try:
+        from ..export import serve_data
+    except ImportError as e:
+        sys.exit(IMPORT_FAILED_MSG.format(e))
+
+    run = open_run(args.proposal, args.run, data='all')
+
+    if not args.include:
+        print("Available sources:")
+        for s in sorted(run.all_sources):
+            print(f"  {s}")
+        sys.exit("Please select at least one source with --include")
+
+    include = []
+    for pat in args.include:
+        if '[' in pat:
+            if not pat.endswith(']'):
+                sys.exit(f"Missing final ] in {pat!r}")
+            src_pat, key_pat = pat[:-1].split('[', 1)
+            include.append((src_pat, key_pat))
+        else:
+            # Source pattern only
+            include.append(pat)
+
+    if args.allow_partial:
+        sel = run.select(include, require_any=True)
+    else:
+        sel = run.select(include, require_all=True)
+
+    try:
+        serve_data(
+            sel, args.port,
+            append_detector_modules=args.append_detector_modules,
+            dummy_timestamps=args.dummy_timestamps,
+            use_infiniband=args.use_infiniband, sock=args.socket_type
+        )
+    except KeyboardInterrupt:
+        print('\nStopped.')
+
+if __name__ == '__main__':
+    main()

extra_data/export.py

@@ -11,6 +11,8 @@
 """
 
 import os.path as osp
+import time
+from collections import deque
 from socket import AF_INET
 from warnings import warn
 
@@ -107,9 +109,7 @@
         yield tid, train_data
 
 
-def serve_files(path, port, source_glob='*', key_glob='*',
-                append_detector_modules=False, dummy_timestamps=False,
-                use_infiniband=False, sock='REP'):
+def serve_files(path, port, source_glob='*', key_glob='*', **kwargs):
     """Stream data from files through a TCP socket.
 
     Parameters
@@ -144,15 +144,58 @@
         data = H5File(path)
 
     data = data.select(source_glob, key_glob)
+    serve_data(data, port, **kwargs)
 
+
+def serve_data(data, port, append_detector_modules=False,
+                dummy_timestamps=False, use_infiniband=False, sock='REP'):
+    """Stream data from files through a TCP socket.
+
+    Parameters
+    ----------
+    data: DataCollection
+        The data to be streamed; should already have sources & keys selected.
+    port: str or int
+        A ZMQ endpoint (e.g. 'tcp://*:44444') or a TCP port to bind the socket
+        to. Integers or strings of all digits are treated as port numbers.
+    append_detector_modules: bool
+        Combine multi-module detector data in a single data source (sources for
+        individual modules are removed). The last section of the source name is
+        replaces with 'APPEND', example:
+            'SPB_DET_AGIPD1M-1/DET/#CH0:xtdf' -> 'SPB_DET_AGIPD1M-1/DET/APPEND'
+
+        Supported detectors: AGIPD, DSSC, LPD
+    dummy_timestamps: bool
+        Whether to add mock timestamps if the metadata lacks them.
+    use_infiniband: bool
+        Use infiniband interface if available (if port specifies a TCP port)
+    sock: str
+        socket type - supported: REP, PUB, PUSH (default REP).
+    """
     if isinstance(port, int) or port.isdigit():
         endpt = f'tcp://{find_infiniband_ip() if use_infiniband else "*"}:{port}'
     else:
         endpt = port
+
     sender = Sender(endpt, sock=sock, dummy_timestamps=dummy_timestamps)
     print(f'Streamer started on: {sender.endpoint}')
+    ntrains = len(data.train_ids)
+
+    sent_times = deque([time.monotonic()], 10)
+    count = 0
+    new_time = 0.
+    def print_update(end='\r'):
+        print(f'Sent {count}/{ntrains} trains - Train ID {tid} - {rate:.1f} Hz', end=end)
+
     for tid, data in _iter_trains(data, merge_detector=append_detector_modules):
         sender.send(data)
+        count += 1
+        new_time = time.monotonic()
+        if count % 5 == 0:
+            rate = len(sent_times) / (new_time - sent_times[0])
+            print_update()
+        sent_times.append(new_time)
+    print_update(end='\n')
 
     # The karabo-bridge code sets linger to 0 so that it doesn't get stuck if
     # the client goes away. But this would also mean that we close the socket

setup.py

@@ -44,6 +44,7 @@ def find_version(*parts):
           "console_scripts": [
               "lsxfel = extra_data.lsxfel:main",
               "karabo-bridge-serve-files = extra_data.cli.serve_files:main",
+              "karabo-bridge-serve-run = extra_data.cli.serve_run:main",
               "extra-data-validate = extra_data.validation:main",
               "extra-data-make-virtual-cxi = extra_data.cli.make_virtual_cxi:main",
               "extra-data-locality = extra_data.locality:main",