Repository: European-XFEL/EXtra-data

Issue:
CLI to stream combined raw+proc data
An instrument scientist spotted that there's no option at present to stream a combination of raw+proc data, even though EXtra-data generally makes it easy to work with both together.

I'd also like to improve selecting sources & keys to stream, and make it harder to accidentally stream all keys from all sources, which tends to be slow. So maybe we can have a new command like `karabo-bridge-serve-run`.

Discussion:

Patch:
diff --git a/.github/dependabot/constraints.txt b/.github/dependabot/constraints.txt
index e1fedc1..a451a4d 100644
--- a/.github/dependabot/constraints.txt
+++ b/.github/dependabot/constraints.txt
@@ -1,16 +1,16 @@
-cloudpickle<=3.0.0
+cloudpickle==2.2.1
 coverage==6.2
-h5py<3.11.0
+h5py<3.10.0
 karabo-bridge==0.7.0
-msgpack<=1.0.7
+msgpack<=1.0.6
 msgpack-numpy==0.4.8
 numpy<1.27.0
 pandas<2.2.0
-Pillow<10.2.0
-psutil==5.9.6
+Pillow<10.1.0
+psutil==5.9.5
 pyparsing==3.1.1
 python-dateutil==2.8.2
 pytz==2023.3.post1
 pyzmq==25.1.1
 six==1.16.0
-xarray<2023.11.0
+xarray<2023.10.0
diff --git a/docs/aligning_trains.ipynb b/docs/aligning_trains.ipynb
index be6605d..d8689bc 100644
--- a/docs/aligning_trains.ipynb
+++ b/docs/aligning_trains.ipynb
@@ -849,7 +849,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "xfel",
+   "display_name": "xfel (Python 3.7)",
    "language": "python",
    "name": "xfel"
   },
diff --git a/docs/cli.rst b/docs/cli.rst
index a62cc47..07fb38b 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -44,6 +44,75 @@ Check the structure of an EuXFEL run or HDF5 file:
 If it finds problems with the data, the program will produce a list of them and
 exit with status 1. See :doc:`validation` for details of what it checks.
 
+.. _cmd-serve-run:
+
+``karabo-bridge-serve-run``
+---------------------------
+
+Stream data from a run, by proposal & run number, in the `Karabo bridge
+<https://rtd.xfel.eu/docs/data-analysis-user-documentation/en/latest/online.html#streaming-from-karabo-bridge>`_
+format. See :doc:`streaming` for more information.
+
+.. code-block:: shell
+
+   #                     Proposal run
+   karabo-bridge-serve-run 700000 40 --port 4321 \
+        --include 'SPB_IRDA_JF4M/DET/JNGFR*:daqOutput' \
+        --include '*/MOTOR/*[*Position]'
+
+.. program:: karabo-bridge-serve-run
+
+.. option:: --port <port>
+
+   Either a numeric TCP port, e.g. ``4321``, or a ZMQ endpoint address, e.g.
+   ``tcp://0.0.0.0:4321``. You will need to give the receiving code this port
+   number or address as well.
+
+   If no port is specified, the program will pick an unused port, and display
+   the endpoint address as it starts.
+
+.. option:: --include <pattern>
+
+   Sources matching the pattern will be included in the streamed data.
+   You can also match keys by putting a key pattern in ``[]`` square brackets
+   at the end of the source pattern.
+
+   You must specify at least one ``--include`` pattern, and you can use the
+   option several times to expand the selection.
+
+   If data is flowing slower than you expect, see if you can use more specific
+   patterns to avoid sending unnecessary sources.
+
+.. option:: --allow-partial
+
+   By default, trains are only sent if they contain all the data selected by
+   ``--include``. This option also sends trains where some of that data is
+   missing.
+
+.. option:: --append-detector-modules
+
+   If the file data contains multiple detector modules as separate sources,
+   i. e. for big area detectors (AGIPD, LPD and DSSC), append these into one
+   single source.
+
+.. option:: --dummy-timestamps
+
+   Add mock timestamps if missing in the original meta-data.
+
+These two options above - appended module sources and dummy timestamps - are
+required if streamed data shall be provided to OnDA.
+
+.. option:: -z <type>, --socket-type <type>
+
+   The ZMQ socket type to use, one of ``PUB``, ``PUSH`` or ``REP``.
+   Default: ``REP``.
+
+.. option:: --use-infiniband
+
+   Use the infiniband network interface (``ib0``) if it's present.
+   This is ignored if ``--port`` is used with a full ZMQ endpoint address.
+
+
 .. _cmd-serve-files:
 
 ``karabo-bridge-serve-files``
@@ -53,6 +122,9 @@ Stream data from files in the `Karabo bridge
 <https://rtd.xfel.eu/docs/data-analysis-user-documentation/en/latest/online.html#streaming-from-karabo-bridge>`_
 format. See :doc:`streaming` for more information.
 
+For streaming data from a run directory, we recommend the newer
+:ref:`cmd-serve-run` command in place of this.
+
 .. code-block:: shell
 
    karabo-bridge-serve-files /gpfs/exfel/exp/XMPL/201750/p700000/proc/r0005 4321
diff --git a/docs/conf.py b/docs/conf.py
index cb8f06c..aead921 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -76,7 +76,7 @@ release = version
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
-language = None
+language = 'en'
 
 # There are two options for replacing |today|: either, you set today to some
 # non-false value, then it is used:
diff --git a/docs/reading_files.rst b/docs/reading_files.rst
index 31b9f63..1939c6f 100644
--- a/docs/reading_files.rst
+++ b/docs/reading_files.rst
@@ -221,10 +221,6 @@ below, e.g.::
    
       .. versionadded:: 1.9
 
-   .. autoattribute:: units
-
-   .. autoattribute:: units_name
-
 
 The run or file object (a :class:`DataCollection`) also has methods to load
 data by sources and keys. :meth:`get_array`, :meth:`get_dask_array` and
diff --git a/docs/streaming.rst b/docs/streaming.rst
index cacfa3e..f0aaea0 100644
--- a/docs/streaming.rst
+++ b/docs/streaming.rst
@@ -6,14 +6,20 @@ socket. The ``extra_data`` Python package can stream data from files using the s
 protocol. You can use this to test code which expects to receive data from
 Karabo Bridge, or use the same code for analysing live data and stored data.
 
-To stream the data from a file or run unmodified, use the command::
+To stream data from a saved run, use the ``karabo-bridge-serve-run`` command:
 
-    karabo-bridge-serve-files /gpfs/exfel/exp/SPB/201830/p900022/raw/r0034 4545
+.. code-block:: shell
 
-The number (4545) must be an unused TCP port above 1024. It will bind to
-this and stream the data to any connected clients.
+   #                     Proposal run
+   karabo-bridge-serve-run 700000 40 --port 4545 \
+        --include 'SPB_IRDA_JF4M/DET/JNGFR*:daqOutput' \
+        --include '*/MOTOR/*[*Position]'
+
+The port number (4545 above) must be an unused TCP port above 1024.
+Clients will need this port and the IP address of the sender to connect.
+For clients running on the same node, use the IP address ``127.0.0.1``.
 Command-line options are explained on the
-:ref:`command reference <cmd-serve-files>` page.
+:ref:`command reference <cmd-serve-run>` page.
 
 .. note::
 
diff --git a/extra_data/cli/serve_run.py b/extra_data/cli/serve_run.py
new file mode 100644
index 0000000..3012295
--- /dev/null
+++ b/extra_data/cli/serve_run.py
@@ -0,0 +1,90 @@
+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", default="0", help="TCP port or ZMQ endpoint to send data on. "
+                                    "Selects a random TCP port by default.")
+    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()
diff --git a/extra_data/export.py b/extra_data/export.py
index 79f6a38..c23c3ed 100644
--- a/extra_data/export.py
+++ b/extra_data/export.py
@@ -11,6 +11,8 @@ program. If not, see <https://opensource.org/licenses/BSD-3-Clause>
 """
 
 import os.path as osp
+import time
+from collections import deque
 from socket import AF_INET
 from warnings import warn
 
@@ -107,9 +109,7 @@ def _iter_trains(data, merge_detector=False):
         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 @@ def serve_files(path, port, source_glob='*', key_glob='*',
         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
+    tid, rate = 0, 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
diff --git a/extra_data/keydata.py b/extra_data/keydata.py
index 3e4430e..b121c11 100644
--- a/extra_data/keydata.py
+++ b/extra_data/keydata.py
@@ -1,13 +1,11 @@
 from typing import List, Optional, Tuple
 
-import h5py
 import numpy as np
 
 from .exceptions import TrainIDError, NoDataError
 from .file_access import FileAccess
 from .read_machinery import (
-    contiguous_regions, DataChunk, select_train_ids, split_trains, roi_shape,
-    trains_files_index,
+    contiguous_regions, DataChunk, select_train_ids, split_trains, roi_shape
 )
 
 class KeyData:
@@ -108,30 +106,6 @@ class KeyData:
         """The size of the data in memory in gigabytes."""
         return self.nbytes / 1e9
 
-    @property
-    def units(self):
-        """The units symbol for this data, e.g. 'μJ', or None if not found"""
-        attrs = self.attributes()
-        base_unit = attrs.get('unitSymbol', None)
-        if base_unit is None:
-            return None
-
-        prefix = attrs.get('metricPrefixSymbol', '')
-        if prefix == 'u':
-            prefix = 'μ'  # We are not afraid of unicode
-        return prefix + base_unit
-
-    @property
-    def units_name(self):
-        """The units name for this data, e.g. 'microjoule', or None if not found"""
-        attrs = self.attributes()
-        base_unit = attrs.get('unitName', None)
-        if base_unit is None:
-            return None
-
-        prefix = attrs.get('metricPrefixName', '')
-        return prefix + base_unit
-
     @property
     def source_file_paths(self):
         paths = []
@@ -155,24 +129,6 @@ class KeyData:
         from pathlib import Path
         return [Path(p) for p in paths]
 
-    def attributes(self):
-        """Get a dict of all attributes stored with this data
-
-        This may be awkward to use. See .units and .units_name for more
-        convenient forms.
-        """
-        dset = self.files[0].file[self.hdf5_data_path]
-        attrs = dict(dset.attrs)
-        if (not attrs) and dset.is_virtual:
-            # Virtual datasets were initially created without these attributes.
-            # Find a source file. Not using source_file_paths as it can give [].
-            _, filename, _, _ = dset.virtual_sources()[0]
-            # Not using FileAccess: no need for train or source lists.
-            with h5py.File(filename, 'r') as f:
-                attrs = dict(f[self.hdf5_data_path].attrs)
-
-        return attrs
-
     def select_trains(self, trains):
         """Select a subset of trains in this data as a new :class:`KeyData` object.
 
@@ -186,16 +142,10 @@ class KeyData:
     def __getitem__(self, item):
         return self.select_trains(item)
 
-    def _only_tids(self, tids, files=None):
+    def _only_tids(self, tids):
         tids_arr = np.array(tids)
-        if files is None:
-            files = [
-                f for f in self.files
-                if f.has_train_ids(tids_arr, self.inc_suspect_trains)
-            ]
-        if not files:
-            # Keep 1 file, even if 0 trains selected.
-            files = [self.files[0]]
+        files = [f for f in self.files
+                 if f.has_train_ids(tids_arr, self.inc_suspect_trains)]
 
         return KeyData(
             self.source,
@@ -236,16 +186,8 @@ class KeyData:
             A maximum number of trains in each part. Parts will often have
             fewer trains than this.
         """
-        # tids_files points to the file for each train.
-        # This avoids checking all files for each chunk, which can be slow.
-        tids_files = trains_files_index(
-            self.train_ids, self.files, self.inc_suspect_trains
-        )
-        for sl in split_trains(len(self.train_ids), parts, trains_per_part):
-            tids = self.train_ids[sl]
-            files = set(tids_files[sl]) - {None}
-            files = sorted(files, key=lambda f: f.filename)
-            yield self._only_tids(tids, files=files)
+        for s in split_trains(len(self.train_ids), parts, trains_per_part):
+            yield self.select_trains(s)
 
     def data_counts(self, labelled=True):
         """Get a count of data entries in each train.
diff --git a/extra_data/read_machinery.py b/extra_data/read_machinery.py
index 65e4279..13e7c6f 100644
--- a/extra_data/read_machinery.py
+++ b/extra_data/read_machinery.py
@@ -4,6 +4,7 @@ The public API is in extra_data.reader; this is internal code.
 """
 import logging
 import math
+import os
 import os.path as osp
 import re
 import time
@@ -17,7 +18,7 @@ log = logging.getLogger(__name__)
 DETECTOR_NAMES = {'AGIPD', 'DSSC', 'LPD'}
 DETECTOR_SOURCE_RE = re.compile(r'(.+)/DET/(\d+)CH')
 
-DATA_ROOT_DIR = '/gpfs/exfel/exp'
+DATA_ROOT_DIR = os.environ.get('EXTRA_DATA_DATA_ROOT', '/gpfs/exfel/exp')
 
 
 class _SliceConstructor(type):
@@ -152,18 +153,6 @@ def split_trains(n_trains, parts=None, trains_per_part=None) -> [slice]:
         for i in range(n_parts)
     ]
 
-def trains_files_index(train_ids, files, inc_suspect_trains=True) -> list:
-    """Make a list of which FileAccess contains each train, used in splitting"""
-    tids_files = [None] * len(train_ids)
-    tid_to_ix = {t: i for i, t in enumerate(train_ids)}
-    for file in files:
-        f_tids = file.train_ids if inc_suspect_trains else file.valid_train_ids
-        for tid in f_tids:
-            ix = tid_to_ix.get(tid, None)
-            if ix is not None:
-                tids_files[ix] = file
-    return tids_files
-
 class DataChunk:
     """Reference to a contiguous chunk of data for one or more trains."""
     def __init__(self, file, dataset_path, first, train_ids, counts):
diff --git a/extra_data/reader.py b/extra_data/reader.py
index 0e19d31..75aa5ac 100644
--- a/extra_data/reader.py
+++ b/extra_data/reader.py
@@ -1013,12 +1013,13 @@ class DataCollection:
                     else:  # require_any
                         train_ids = np.union1d(train_ids, source_tids)
 
-            train_ids = list(train_ids)  # Convert back to a list.
             sources_data = {
                 src: srcdata._only_tids(train_ids)
                 for src, srcdata in sources_data.items()
             }
 
+            train_ids = list(train_ids)  # Convert back to a list.
+
         else:
             train_ids = self.train_ids
 
@@ -1122,28 +1123,8 @@ class DataCollection:
             A maximum number of trains in each part. Parts will often have
             fewer trains than this.
         """
-        for source in self._sources_data.values():
-            assert source.train_ids == self.train_ids
-
-        def dict_zip(iter_d):
-            while True:
-                try:
-                    yield {k: next(v) for (k, v) in iter_d.items()}
-                except StopIteration:
-                    return
-
-        for sources_data_part in dict_zip({
-            n: s.split_trains(parts=parts, trains_per_part=trains_per_part)
-            for (n, s) in self._sources_data.items()
-        }):
-            files = set().union(*[sd.files for sd in sources_data_part.values()])
-            train_ids = list(sources_data_part.values())[0].train_ids
-
-            yield DataCollection(
-                files, sources_data=sources_data_part, train_ids=train_ids,
-                aliases=self._aliases, inc_suspect_trains=self.inc_suspect_trains,
-                is_single_run=self.is_single_run,
-            )
+        for s in split_trains(len(self.train_ids), parts, trains_per_part):
+            yield self.select_trains(s)
 
     def _check_source_conflicts(self):
         """Check for data with the same source and train ID in different files.
@@ -1216,6 +1197,14 @@ class DataCollection:
             span_txt = '{}.{}'.format(datetime.timedelta(seconds=seconds),
                                       int(deciseconds))
 
+        detector_modules = {}
+        for source in self.detector_sources:
+            name, modno = DETECTOR_SOURCE_RE.match(source).groups((1, 2))
+            detector_modules[(name, modno)] = source
+
+        # A run should only have one detector, but if that changes, don't hide it
+        detector_name = ','.join(sorted(set(k[0] for k in detector_modules)))
+
         # disp
         print('# of trains:   ', train_count)
         print('Duration:      ', span_txt)
@@ -1223,34 +1212,22 @@ class DataCollection:
         print('Last train ID: ', last_train)
         print()
 
-        if not details_for_sources:
-            # Include summary section for multi-module detectors unless
-            # source details are enabled.
-
-            detector_modules = {}
-            for source in self.detector_sources:
-                name, modno = DETECTOR_SOURCE_RE.match(source).groups((1, 2))
-                detector_modules[(name, modno)] = source
-
-            # A run should only have one detector, but if that changes, don't hide it
-            detector_name = ','.join(sorted(set(k[0] for k in detector_modules)))
-
-            print("{} XTDF detector modules ({})".format(
-                len(self.detector_sources), detector_name
+        print("{} detector modules ({})".format(
+            len(self.detector_sources), detector_name
+        ))
+        if len(detector_modules) > 0:
+            # Show detail on the first module (the others should be similar)
+            mod_key = sorted(detector_modules)[0]
+            mod_source = detector_modules[mod_key]
+            dinfo = self.detector_info(mod_source)
+            module = ' '.join(mod_key)
+            dims = ' x '.join(str(d) for d in dinfo['dims'])
+            print("  e.g. module {} : {} pixels".format(module, dims))
+            print("  {}".format(mod_source))
+            print("  {} frames per train, up to {} frames total".format(
+                dinfo['frames_per_train'], dinfo['total_frames']
             ))
-            if len(detector_modules) > 0:
-                # Show detail on the first module (the others should be similar)
-                mod_key = sorted(detector_modules)[0]
-                mod_source = detector_modules[mod_key]
-                dinfo = self.detector_info(mod_source)
-                module = ' '.join(mod_key)
-                dims = ' x '.join(str(d) for d in dinfo['dims'])
-                print("  e.g. module {} : {} pixels".format(module, dims))
-                print("  {}".format(mod_source))
-                print("  {} frames per train, up to {} frames total".format(
-                    dinfo['frames_per_train'], dinfo['total_frames']
-                ))
-            print()
+        print()
 
         # Invert aliases for faster lookup.
         src_aliases = defaultdict(set)
@@ -1297,16 +1274,9 @@ class DataCollection:
 
                 print(f"{prefix}{k}{alias_str}\t[{dt}{entry_info}]")
 
-        if details_for_sources:
-            # All instrument sources with details enabled.
-            displayed_inst_srcs = self.instrument_sources
-            print(len(displayed_inst_srcs), 'instrument sources:')
-        else:
-            # Only non-XTDF instrument sources without details enabled.
-            displayed_inst_srcs = self.instrument_sources - self.detector_sources
-            print(len(displayed_inst_srcs), 'instrument sources (excluding XTDF detectors):')
-
-        for s in sorted(displayed_inst_srcs):
+        non_detector_inst_srcs = self.instrument_sources - self.detector_sources
+        print(len(non_detector_inst_srcs), 'instrument sources (excluding detectors):')
+        for s in sorted(non_detector_inst_srcs):
             print('  -', s, src_alias_list(s))
             if not any(p.match(s) for p in details_sources_re):
                 continue
diff --git a/extra_data/sourcedata.py b/extra_data/sourcedata.py
index 3a7add7..e130ec1 100644
--- a/extra_data/sourcedata.py
+++ b/extra_data/sourcedata.py
@@ -8,9 +8,7 @@ import h5py
 from .exceptions import MultiRunError, PropertyNameError, NoDataError
 from .file_access import FileAccess
 from .keydata import KeyData
-from .read_machinery import (
-    glob_wildcards_re, same_run, select_train_ids, split_trains, trains_files_index
-)
+from .read_machinery import glob_wildcards_re, same_run, select_train_ids, split_trains
 
 
 class SourceData:
@@ -264,21 +262,16 @@ class SourceData:
         """
         return self._only_tids(select_train_ids(self.train_ids, trains))
 
-    def _only_tids(self, tids, files=None) -> 'SourceData':
-        if files is None:
-            files = [
-                f for f in self.files
-                if f.has_train_ids(tids, self.inc_suspect_trains)
-            ]
-        if not files:
-            # Keep 1 file, even if 0 trains selected, to get keys, dtypes, etc.
-            files = [self.files[0]]
-
+    def _only_tids(self, tids) -> 'SourceData':
         return SourceData(
             self.source,
             sel_keys=self.sel_keys,
             train_ids=tids,
-            files=files,
+            # Keep 1 file, even if 0 trains selected, to get keys, dtypes, etc.
+            files=[
+                f for f in self.files
+                if f.has_train_ids(tids, self.inc_suspect_trains)
+            ] or [self.files[0]],
             section=self.section,
             is_single_run=self.is_single_run,
             inc_suspect_trains=self.inc_suspect_trains
@@ -316,16 +309,8 @@ class SourceData:
             A maximum number of trains in each part. Parts will often have
             fewer trains than this.
         """
-        # tids_files points to the file for each train.
-        # This avoids checking all files for each chunk, which can be slow.
-        tids_files = trains_files_index(
-            self.train_ids, self.files, self.inc_suspect_trains
-        )
-        for sl in split_trains(len(self.train_ids), parts, trains_per_part):
-            tids = self.train_ids[sl]
-            files = set(tids_files[sl]) - {None}
-            files = sorted(files, key=lambda f: f.filename)
-            yield self._only_tids(tids, files=files)
+        for s in split_trains(len(self.train_ids), parts, trains_per_part):
+            yield self.select_trains(s)
 
     def data_counts(self, labelled=True, index_group=None):
         """Get a count of data entries in each train.
diff --git a/setup.py b/setup.py
index 0eaafd9..889e008 100755
--- a/setup.py
+++ b/setup.py
@@ -44,6 +44,7 @@ setup(name="EXtra-data",
           "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",


Test patch:
diff --git a/extra_data/tests/conftest.py b/extra_data/tests/conftest.py
index 9801f3b..abc9530 100644
--- a/extra_data/tests/conftest.py
+++ b/extra_data/tests/conftest.py
@@ -140,7 +140,7 @@ def mock_spb_raw_run(format_version):
         yield td
 
 
-@pytest.fixture(scope='session')
+@pytest.fixture()
 def mock_spb_raw_and_proc_run():
     with TemporaryDirectory() as td:
         prop_dir = osp.join(str(td), 'SPB', '201830', 'p002012')
diff --git a/extra_data/tests/mockdata/xgm.py b/extra_data/tests/mockdata/xgm.py
index d25f737..5f24747 100644
--- a/extra_data/tests/mockdata/xgm.py
+++ b/extra_data/tests/mockdata/xgm.py
@@ -1,5 +1,3 @@
-import numpy as np
-
 from .base import DeviceBase
 
 class XGM(DeviceBase):
@@ -55,16 +53,3 @@ class XGM(DeviceBase):
         ('xTD', 'f4', (1000,)),
         ('yTD', 'f4', (1000,)),
     ]
-
-    def write_instrument(self, f):
-        super().write_instrument(f)
-
-        # Annotate intensityTD with some units to test retrieving them
-        # Karabo stores ASCII strings, assigning bytes is a shortcut to mimic that
-        ds = f[f'INSTRUMENT/{self.device_id}:output/data/intensityTD']
-        ds.attrs['metricPrefixEnum']= np.array([14], dtype=np.int32)
-        ds.attrs['metricPrefixName'] = b'micro'
-        ds.attrs['metricPrefixSymbol'] = b'u'
-        ds.attrs['unitEnum'] = np.array([15], dtype=np.int32)
-        ds.attrs['unitName'] = b'joule'
-        ds.attrs['unitSymbol'] = b'J'
diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py
index c9f00f8..d2a9147 100644
--- a/extra_data/tests/test_keydata.py
+++ b/extra_data/tests/test_keydata.py
@@ -49,7 +49,7 @@ def test_select_trains(mock_spb_raw_run):
     # Empty selection
     sel2 = xgm_beam_x[80:]
     assert sel2.shape == (0,)
-    assert len(sel2.files) == 1
+    assert len(sel2.files) == 0
     assert sel2.xarray().shape == (0,)
 
     # Single train
@@ -339,24 +339,3 @@ def test_file_no_trains(run_with_file_no_trains):
     run = RunDirectory(run_with_file_no_trains)
     xpos = run['SPB_XTD9_XGM/DOOCS/MAIN', 'beamPosition.ixPos'].ndarray()
     assert xpos.shape == (64,)
-
-
-def test_attributes(mock_sa3_control_data):
-    run = H5File(mock_sa3_control_data)
-    xgm_intensity = run['SA3_XTD10_XGM/XGM/DOOCS:output', 'data.intensityTD']
-    attrs = xgm_intensity.attributes()
-
-    assert isinstance(attrs, dict)
-    assert attrs['metricPrefixName'] == 'micro'
-    assert attrs['unitSymbol'] == 'J'
-
-
-def test_units(mock_sa3_control_data):
-    run = H5File(mock_sa3_control_data)
-    xgm_intensity = run['SA3_XTD10_XGM/XGM/DOOCS:output', 'data.intensityTD']
-
-    assert xgm_intensity.units == 'μJ'
-    assert xgm_intensity.units_name == 'microjoule'
-
-    # Check that it still works after selecting 0 trains
-    assert xgm_intensity.select_trains(np.s_[:0]).units == 'μJ'
diff --git a/extra_data/tests/test_reader_mockdata.py b/extra_data/tests/test_reader_mockdata.py
index 846808e..e19f87d 100644
--- a/extra_data/tests/test_reader_mockdata.py
+++ b/extra_data/tests/test_reader_mockdata.py
@@ -702,8 +702,6 @@ def test_split_trains(mock_fxe_raw_run):
     chunks = list(run.split_trains(3))
     assert len(chunks) == 3
     assert {len(c.train_ids) for c in chunks} == {160}
-    arr = chunks[0]['FXE_XAD_GEC/CAM/CAMERA:daqOutput', 'data.image.dims'].ndarray()
-    assert arr.shape == (160, 2)
 
     chunks = list(run.split_trains(4, trains_per_part=100))
     assert len(chunks) == 5
diff --git a/extra_data/tests/test_sourcedata.py b/extra_data/tests/test_sourcedata.py
index a299631..2992afe 100644
--- a/extra_data/tests/test_sourcedata.py
+++ b/extra_data/tests/test_sourcedata.py
@@ -111,8 +111,6 @@ def test_split_trains(mock_spb_raw_run):
     chunks = list(xgm.split_trains(3))
     assert len(chunks) == 3
     assert {len(c.train_ids) for c in chunks} == {21, 22}
-    # The middle chunk spans across 2 files
-    assert [len(c.files) for c in chunks] == [1, 2, 1]
 
     chunks = list(xgm.split_trains(3, trains_per_part=20))
     assert len(chunks) == 4
diff --git a/extra_data/tests/test_streamer.py b/extra_data/tests/test_streamer.py
index 3062e93..9aa84e2 100644
--- a/extra_data/tests/test_streamer.py
+++ b/extra_data/tests/test_streamer.py
@@ -1,9 +1,11 @@
 """Test streaming data with ZMQ interface."""
 
 import os
+import signal
+from subprocess import PIPE, Popen, TimeoutExpired
 
+import numpy as np
 import pytest
-from subprocess import PIPE, Popen
 
 from extra_data import by_id, H5File, RunDirectory
 from extra_data.export import _iter_trains, ZMQStreamer
@@ -41,6 +43,19 @@ def test_merge_detector(mock_fxe_raw_run, mock_fxe_control_data, mock_spb_proc_r
             break
 
 
+def cleanup_proc(p: Popen):
+    if p.poll() is None:
+        p.send_signal(signal.SIGINT)
+        try:
+            p.wait(timeout=2)
+        except TimeoutExpired:
+            pass
+    if p.poll() is None:
+        p.kill()
+        rc = p.wait(timeout=2)
+        assert rc == -9  # process terminated by kill signal
+
+
 @pytest.mark.skipif(os.name != 'posix', reason="Test uses Unix socket")
 def test_serve_files(mock_fxe_raw_run, tmp_path):
     src = 'FXE_XAD_GEC/CAM/CAMERA:daqOutput'
@@ -67,10 +82,38 @@ def test_serve_files(mock_fxe_raw_run, tmp_path):
         assert tid == 10000
         assert set(data) == {src}
     finally:
-        if p.poll() is None:
-            p.kill()
-            rc = p.wait(timeout=2)
-            assert rc == -9  # process terminated by kill signal
+        cleanup_proc(p)
+
+
+@pytest.mark.skipif(os.name != 'posix', reason="Test uses Unix socket")
+def test_serve_run(mock_spb_raw_and_proc_run, tmp_path):
+    mock_data_root, _, _ = mock_spb_raw_and_proc_run
+    zmq_endpoint = f'ipc://{tmp_path}/socket'
+    xgm_src = 'SPB_XTD9_XGM/DOOCS/MAIN'
+    agipd_m0_src = 'SPB_DET_AGIPD1M-1/DET/0CH0:xtdf'
+    args = ['karabo-bridge-serve-run', '2012', '238',
+            '--port', zmq_endpoint,
+            '--include', f'{xgm_src}[beamPosition.i*Pos]',
+            '--include', '*AGIPD1M-1/DET/0CH0:xtdf'
+           ]
+
+    p = Popen(args, env=dict(
+        os.environ,
+        PYTHONUNBUFFERED='1',
+        EXTRA_DATA_DATA_ROOT=mock_data_root
+    ))
+    try:
+        with Client(zmq_endpoint, timeout=30) as c:
+            data, meta = c.next()
+
+        tid = next(m['timestamp.tid'] for m in meta.values())
+        assert tid == 10000
+        assert set(data) == {xgm_src, agipd_m0_src}
+        assert set(data[xgm_src]) == \
+               {f'beamPosition.i{xy}Pos.value' for xy in 'xy'} | {'metadata'}
+        assert data[agipd_m0_src]['image.data'].dtype == np.float32
+    finally:
+        cleanup_proc(p)
 
 
 def test_deprecated_server():
diff --git a/extra_data/tests/test_voview.py b/extra_data/tests/test_voview.py
index 60df0d3..cbdfbe1 100644
--- a/extra_data/tests/test_voview.py
+++ b/extra_data/tests/test_voview.py
@@ -63,8 +63,6 @@ def test_use_voview(mock_spb_raw_run, tmp_path):
     assert {p.name for p in xgm_intens[:30].source_file_paths} == {
         'RAW-R0238-DA01-S00000.h5'
     }
-    assert xgm_intens.units == 'μJ'
-    assert xgm_intens.units_name == 'microjoule'