mccode-plumber 0.6.0__tar.gz

mccode_plumber-0.6.0/.github/workflows/pip.yml

@@ -0,0 +1,38 @@
+name: Pip
+
+on:
+  workflow_dispatch:
+    #  pull_request:
+    #  push:
+    #    branches:
+    #      - main
+
+jobs:
+  build:
+    name: Build with Pip
+    runs-on: ${{ matrix.platform }}
+    strategy:
+      fail-fast: false
+      matrix:
+        platform: [windows-latest, macos-latest, ubuntu-latest]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Set min macOS version
+      if: runner.os == 'macOS'
+      run: |
+        echo "MACOS_DEPLOYMENT_TARGET=10.14" >> $GITHUB_ENV
+
+    - name: Build and install
+      run: pip install --verbose .
+
+    - name: Test
+      run: |
+        python -m pip install pytest
+        python -m pytest

mccode_plumber-0.6.0/.github/workflows/wheels.yml

@@ -0,0 +1,52 @@
+name: Wheels
+
+on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches:
+     - main
+  release:
+    types:
+      - published
+
+jobs:
+  build_sdist_and_wheel:
+    name: Build SDist and Wheel
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+        submodules: true
+
+    - name: Build SDist
+      run: pipx run build
+
+    - name: Check metadata
+      run: pipx run twine check dist/*
+
+    - uses: actions/upload-artifact@v4
+      with:
+        path: dist/*
+
+  upload_all:
+    name: Upload if release
+    needs: [build_sdist_and_wheel]
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/mccode-plumber
+    permissions:
+      id-token: write
+    if: github.event_name == 'release' && github.event.action == 'published'
+
+    steps:
+    - uses: actions/setup-python@v5
+
+    - uses: actions/download-artifact@v4
+      with:
+        name: artifact
+        path: dist
+
+    - uses: pypa/gh-action-pypi-publish@release/v1

mccode_plumber-0.6.0/.gitignore

@@ -0,0 +1,5 @@
+build/
+venv/
+__pycache__/
+.idea/
+*.egg-info

mccode_plumber-0.6.0/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: mccode-plumber
+Version: 0.6.0
+Author-email: Gregory Tucker <gregory.tucker@ess.eu>
+Classifier: License :: OSI Approved :: BSD License
+Description-Content-Type: text/markdown
+Requires-Dist: p4p
+Requires-Dist: file-writer-control>=1.3.0
+Requires-Dist: restage>=0.4.0
+Requires-Dist: mccode-to-kafka>=0.2.1
+Requires-Dist: moreniius>=0.2.3
+Requires-Dist: icecream
+
+# McCode Plumber
+Setup, run, and teardown the infrastructure for splitrun McCode scans sending data through Kafka into NeXus

mccode_plumber-0.6.0/README.md

@@ -0,0 +1,2 @@
+# McCode Plumber
+Setup, run, and teardown the infrastructure for splitrun McCode scans sending data through Kafka into NeXus

mccode_plumber-0.6.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ['setuptools>=60', 'setuptools_scm>=8.0']
+build-backend = 'setuptools.build_meta'
+
+[project]
+name = 'mccode-plumber'
+dependencies = [
+    'p4p',
+    'file-writer-control>=1.3.0',
+    'restage>=0.4.0',
+    'mccode-to-kafka>=0.2.1',
+    'moreniius>=0.2.3',
+    'icecream',
+]
+readme = "README.md"
+authors = [
+    { name = "Gregory Tucker", email = "gregory.tucker@ess.eu" },
+]
+classifiers = [
+    "License :: OSI Approved :: BSD License",
+]
+dynamic = ['version']
+
+[project.scripts]
+mp-splitrun = 'mccode_plumber.splitrun:main'
+mp-epics = 'mccode_plumber.epics:run'
+mp-epics-update = 'mccode_plumber.epics:update'
+mp-forwarder-setup = 'mccode_plumber.forwarder:setup'
+mp-forwarder-teardown = 'mccode_plumber.forwarder:teardown'
+mp-writer-from = 'mccode_plumber.writer:print_time'
+mp-writer-write = 'mccode_plumber.writer:start_writer'
+mp-writer-wait = 'mccode_plumber.writer:wait_on_writer'
+mp-register-topics = 'mccode_plumber.kafka:register_topics'
+mp-insert-hdf5-instr = 'mccode_plumber.mccode:insert'
+
+[tool.setuptools_scm]
+

mccode_plumber-0.6.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mccode_plumber-0.6.0/src/mccode_plumber/epics.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+from p4p.nt import NTScalar
+from p4p.server import Server, StaticProvider
+from p4p.server.thread import SharedPV
+from pathlib import Path
+from typing import Union
+
+
+def convert_instr_parameters_to_nt(parameters):
+    from mccode_antlr.common.expression import DataType, ShapeType
+    out = {}
+    for p in parameters:
+        expr = p.value
+        if expr.is_str:
+            t, d = 's', ''
+        elif expr.data_type == DataType.int:
+            t, d = 'i', 0
+        elif expr.data_type == DataType.float:
+            t, d = 'd', 0.0
+        else:
+            raise ValueError(f"Unknown parameter type {expr.data_type}")
+        if expr.shape_type == ShapeType.vector:
+            t, d = 'a' + t, [d]
+        out[p.name] = NTScalar(t).wrap(expr.value if expr.has_value else d)
+    return out
+
+
+def parse_instr_nt_values(instr: Union[Path, str]):
+    """Get the instrument parameters from an Instr a or a parseable Instr file and convert to NTScalar values"""
+    from .mccode import get_mccode_instr_parameters
+    nts = convert_instr_parameters_to_nt(get_mccode_instr_parameters(instr))
+    if 'mcpl_filename' not in nts:
+        nts['mcpl_filename'] = NTScalar('s').wrap('')
+    return nts
+
+
+class MailboxHandler:
+    @staticmethod
+    def put(pv, op):
+        from datetime import datetime, timezone
+        val = op.value()
+
+        if pv.nt is None:
+            # Assume that this means wrap wasn't provided ...
+            pv.nt = NTScalar(val.type()['value'])
+            pv._wrap = pv.nt.wrap
+
+        # Notify any subscribers of the new value, adding the timestamp, so they know when it was set.
+        pv.post(val, timestamp=datetime.now(timezone.utc).timestamp())
+        # Notify the client making this PUT operation that it has now completed
+        op.done()
+
+
+def get_parser():
+    from argparse import ArgumentParser
+    p = ArgumentParser()
+    p.add_argument('instr', type=str, help='The instrument file to read')
+    p.add_argument('-p', '--prefix', type=str, help='The EPICS PV prefix to use', default='mcstas:')
+    return p
+
+
+def parse_args():
+    args = get_parser().parse_args()
+    parameters = parse_instr_nt_values(args.instr)
+    return parameters, args
+
+
+def main(names: dict[str, NTScalar], prefix: str = None):
+    provider = StaticProvider('mailbox')  # 'mailbox' is an arbitrary name
+
+    pvs = []  # we must keep a reference in order to keep the Handler from being collected
+    for name, value in names.items():
+        pv = SharedPV(initial=value, handler=MailboxHandler())
+        provider.add(f'{prefix}{name}' if prefix else name, pv)
+        pvs.append(pv)
+
+    print(f'Start mailbox server for {len(pvs)} PVs with prefix {prefix}')
+    Server.forever(providers=[provider])
+    print('Done')
+
+
+def run():
+    parameters, args = parse_args()
+    main(parameters, prefix=args.prefix)
+
+
+def start(parameters, prefix: str = None):
+    from multiprocessing import Process
+    proc = Process(target=main, args=(parameters, prefix))
+    proc.start()
+    return proc
+
+
+def stop(proc):
+    proc.terminate()
+    proc.join(1)
+    proc.close()
+
+
+def update():
+    from argparse import ArgumentParser
+    from p4p.client.thread import Context
+    parser = ArgumentParser(description="Update the mailbox server with new values")
+    parser.add_argument('address', type=str, help='The mailbox address of the value to be updated')
+    parser.add_argument('value', type=str, help='The new value to be assigned to the mailbox')
+    args = parser.parse_args()
+
+    ctx = Context('pva')
+    pv = ctx.get(args.address, throw=False)
+    if isinstance(pv, float):
+        ctx.put(args.address, float(args.value))
+    elif isinstance(pv, int):
+        ctx.put(args.address, int(args.value))
+    elif isinstance(pv, str):
+        ctx.put(args.address, str(args.value))
+    else:
+        raise ValueError(f'Unknown type {type(pv)} (this is likely a vector that I can not handle yet?)')
+
+    ctx.disconnect()
+
+
+if __name__ == '__main__':
+    run()

mccode_plumber-0.6.0/src/mccode_plumber/forwarder.py

@@ -0,0 +1,88 @@
+"""
+Control a running Forwarder instance to send data to a Kafka broker.
+
+Two gateway functions are exposed as system scripts to add and remove an Instr's parameters from the Forwarder's
+list of EPICS PVs to monitor.
+
+Alternatively, the same functionality can be accessed from Python using the configure_forwarder and reset_forwarder
+functions. Which take PV information and Forwarder/Kafka configuration as arguments.
+"""
+
+
+def normalise_pvs(pvs: list[dict], config=None, prefix=None, topic=None):
+    if config is None:
+        config = "localhost:9092/forwarderConfig"
+    if prefix is None:
+        prefix = 'mcstas:'
+    if topic is None:
+        topic = 'mcstasParameters'
+
+    if '/' not in config:
+        raise RuntimeError('Expected / to separate broker and topic in Forwarder Kafka configuration specification')
+
+    cfg_broker, cfg_topic = config.split('/', 1)
+
+    for pv in pvs:
+        if 'source' not in pv:
+            pv['source'] = f'{prefix}{pv["name"]}'
+        if 'topic' not in pv:
+            pv['topic'] = topic
+    return cfg_broker, cfg_topic, pvs
+
+
+def streams(pvs: list[dict]):
+    from streaming_data_types.forwarder_config_update_rf5k import StreamInfo, Protocol
+    return [StreamInfo(pv['source'], pv['module'], pv['topic'], Protocol.Protocol.PVA) for pv in pvs]
+
+
+def configure_forwarder(pvs: list[dict], config=None, prefix=None, topic=None):
+    from confluent_kafka import Producer
+    from streaming_data_types.forwarder_config_update_rf5k import serialise_rf5k, StreamInfo, Protocol
+    from streaming_data_types.fbschemas.forwarder_config_update_rf5k.UpdateType import UpdateType
+
+    cfg_broker, cfg_topic, pvs = normalise_pvs(pvs, config, prefix, topic)
+    producer = Producer({"bootstrap.servers": cfg_broker})
+    producer.produce(cfg_topic, serialise_rf5k(UpdateType.ADD, streams(pvs)))
+    producer.flush()
+    return pvs
+
+
+def reset_forwarder(pvs: list[dict], config=None, prefix=None, topic=None):
+    from confluent_kafka import Producer
+    from streaming_data_types.forwarder_config_update_rf5k import serialise_rf5k
+    from streaming_data_types.fbschemas.forwarder_config_update_rf5k.UpdateType import UpdateType
+
+    cfg_broker, cfg_topic, pvs = normalise_pvs(pvs, config, prefix, topic)
+    producer = Producer({"bootstrap.servers": cfg_broker})
+    producer.produce(cfg_topic, serialise_rf5k(UpdateType.REMOVE, streams(pvs)))
+    producer.flush()
+    return pvs
+
+
+def parse_registrar_args():
+    from argparse import ArgumentParser
+    from .mccode import get_mccode_instr_parameters
+
+    parser = ArgumentParser(description="Discover EPICS PVs and inform a forwarder about them")
+    parser.add_argument('-p', '--prefix', type=str, default='mcstas:')
+    parser.add_argument('instrument', type=str, help="The mcstas instrument with EPICS PVs")
+    parser.add_argument('-c', '--config', type=str, help="The Kafka server and topic for configuring the Forwarder")
+    parser.add_argument('-t', '--topic', type=str, help="The Kafka topic to instruct the Forwarder to send data to")
+
+    args = parser.parse_args()
+    parameter_names = [p.name for p in get_mccode_instr_parameters(args.instrument)]
+    if 'mcpl_filename' not in parameter_names:
+        parameter_names.append('mcpl_filename')
+    # the forwarder only cares about: "source", "module", "topic"
+    params = [{'source': f'{args.prefix}{name}', 'module': 'f144', 'topic': args.topic} for name in parameter_names]
+    return params, args
+
+
+def setup():
+    parameters, args = parse_registrar_args()
+    configure_forwarder(parameters, config=args.config, prefix=args.prefix, topic=args.topic)
+
+
+def teardown():
+    parameters, args = parse_registrar_args()
+    reset_forwarder(parameters, config=args.config, prefix=args.prefix, topic=args.topic)

mccode_plumber-0.6.0/src/mccode_plumber/kafka.py

@@ -0,0 +1,27 @@
+def parse_kafka_topic_args():
+    from argparse import ArgumentParser
+    parser = ArgumentParser(description="Prepare the named Kafka broker to host one or more topics")
+    parser.add_argument('-b', '--broker', type=str, help='The Kafka broker server to interact with')
+    parser.add_argument('topic', nargs="+", type=str, help='The Kafka topic(s) to register')
+    parser.add_argument('-q', '--quiet', action='store_true', help='Quiet (positive) failure')
+
+    args = parser.parse_args()
+    return args
+
+
+def register_topics():
+    from confluent_kafka.admin import AdminClient, NewTopic
+    args = parse_kafka_topic_args()
+
+    client = AdminClient({"bootstrap.servers": args.broker})
+    topics = [NewTopic(t, num_partitions=1, replication_factor=1) for t in args.topic]
+    futures = client.create_topics(topics)
+
+    for topic, future in futures.items():
+        try:
+            future.result()
+            print(f"Topic {topic} created")
+        except Exception as e:
+            from confluent_kafka.error import KafkaError
+            if not (args.quiet and e.args[0] == KafkaError.TOPIC_ALREADY_EXISTS):
+                print(f"Failed to create topic {topic}: {e.args[0].str()}")

mccode_plumber-0.6.0/src/mccode_plumber/mccode.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+from typing import Union
+from mccode_antlr.instr import Instr
+from mccode_antlr.common import InstrumentParameter
+
+
+def get_mcstas_instr(filename: Union[Path, str]) -> Instr:
+    from restage.instr import load_instr
+    return load_instr(filename)
+
+
+def get_mccode_instr_parameters(filename: Union[Path, str]) -> tuple[InstrumentParameter]:
+    from mccode_antlr.loader.loader import parse_mccode_instr_parameters
+    if not isinstance(filename, Path):
+        filename = Path(filename)
+    if filename.suffix == '.instr':
+        with filename.open('r') as file:
+            contents = file.read()
+        return parse_mccode_instr_parameters(contents)
+    # otherwise:
+    return get_mcstas_instr(filename).parameters
+
+
+def insert_mcstas_hdf5(filename: Union[Path, str], outfile: Union[Path, str], parent: str):
+    import h5py
+    from mccode_antlr.io.hdf5 import HDF5IO
+    if isinstance(filename, str):
+        filename = Path(filename)
+    with h5py.File(outfile, mode='r+') as dest_file:
+        if parent in dest_file:
+            raise RuntimeError(f'{outfile} already contains an object named {parent}')
+        dest = dest_file.create_group(parent)
+        if filename.stem.lower() in ('.h5', '.hdf', '.hdf5', ):
+            # copy the file contents if it _is_ a serialized instrument
+            with h5py.File(filename, mode='r') as source:
+                for obj in source.keys():
+                    source.copy(obj, dest)
+        else:
+            instr = get_mcstas_instr(filename)
+            HDF5IO.save(dest, instr)
+
+
+def get_arg_parser():
+    from argparse import ArgumentParser
+    from .utils import is_readable, is_appendable
+    parser = ArgumentParser(description="Copy a Instr HDF5 representation to a NeXus HDF5 file")
+    a = parser.add_argument
+    a('instrument', type=is_readable, default=None, help="The mcstas instrument file")
+    a('-p', '--parent', type=str, default='mcstas')
+    a('-o', '--outfile', type=is_appendable, default=None, help='Base NeXus structure, will be extended')
+    return parser
+
+
+def insert():
+    parser = get_arg_parser()
+    args = parser.parse_args()
+    insert_mcstas_hdf5(args.instrument, args.outfile, args.parent)
+
+

mccode_plumber-0.6.0/src/mccode_plumber/splitrun.py

@@ -0,0 +1,28 @@
+def make_parser():
+    from restage.splitrun import make_splitrun_parser
+    parser = make_splitrun_parser()
+    parser.prog = 'mp-splitrun'
+    parser.add_argument('--broker', type=str, help='The Kafka broker to send monitors to', default=None)
+    parser.add_argument('--source', type=str, help='The Kafka source name to use for monitors', default=None)
+    return parser
+
+
+def monitors_to_kafka_callback_with_arguments(broker: str, source: str):
+    from functools import partial
+    from mccode_to_kafka.sender import send_histograms
+
+    def callback(*args, **kwargs):
+        print(f'monitors to kafka callback called with {args} and {kwargs}')
+        return send_histograms(*args, broker=broker, source=source, **kwargs)
+
+    # return partial(send_histograms, broker=broker, source=source), {'dir': 'root'}
+    return callback, {'dir': 'root'}
+
+
+def main():
+    from .mccode import get_mcstas_instr
+    from restage.splitrun import splitrun_args, parse_splitrun
+    args, parameters, precision = parse_splitrun(make_parser())
+    instr = get_mcstas_instr(args.instrument[0])
+    callback, callback_args = monitors_to_kafka_callback_with_arguments(args.broker, args.source)
+    return splitrun_args(instr, parameters, precision, args, callback=callback, callback_arguments=callback_args)

mccode_plumber-0.6.0/src/mccode_plumber/utils.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+
+def is_accessible(access_type):
+    def checker(name: str | None | Path):
+        if name is None or name == '':
+            return None
+        from os import access
+        if not isinstance(name, Path):
+            name = Path(name).resolve()
+        if not name.exists():
+            raise RuntimeError(f'The specified filename {name} does not exist')
+        if not access(name, access_type):
+            raise RuntimeError(f'The specified filename {name} is not {access_type}')
+        return name
+
+    return checker
+
+
+def is_readable(value: str | None | Path):
+    from os import R_OK as READABLE
+    return is_accessible(READABLE)(value)
+
+
+def is_writable(value: str | None | Path):
+    """Determine if a provided path represents an existing writable file"""
+    from os import W_OK
+    if value is None or value == '':
+        return None
+    if not isinstance(value, Path):
+        value = Path(value).resolve()
+    # Typically we can create a new file if the containing folder is writable
+    if not value.exists():
+        return value if is_accessible(W_OK)(value.parent) else None
+    # And if the file exists we should check if we can overwrite it
+    return is_accessible(W_OK)(value)
+
+
+def is_creatable(value: str | None | Path):
+    """Determine if a provided path represents a file that can be created"""
+    if value is None or value == '':
+        return None
+    if not isinstance(value, Path):
+        value = Path(value).resolve()
+    if value.exists():
+        raise RuntimeError(f"The specified filename {value} already exists!")
+    return value if is_writable(value.parent) else None
+
+
+def is_appendable(value: str | None | Path):
+    from os import W_OK, R_OK
+    return is_accessible(R_OK | W_OK)(value)
+
+
+def is_executable(value: str | None | Path):
+    from os import X_OK as EXECUTABLE
+    return is_accessible(EXECUTABLE)(value)
+
+
+def is_callable(name: str | None):
+    if name is None:
+        return None
+    from importlib import import_module
+    module_name, func_name = name.split(':')
+    module = import_module(module_name)
+    return getattr(module, func_name)

mccode_plumber-0.6.0/src/mccode_plumber/writer.py

@@ -0,0 +1,336 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Union, Callable
+from mccode_antlr.instr import Instr
+
+
+def _is_group(x, group):
+    """Is a (dict) object a (NeXus) group with the specified name?"""
+    return 'name' in x and 'type' in x and x['type'] == 'group' and x['name'] == group
+
+
+def _is_stream(x, name):
+    """Is a (dict) object a stream with the specified name?"""
+    return 'module' in x and x['module'] == name
+
+
+def _make_group(name: str, nx_class: str):
+    """Make a (NeXus) group dict with the specified name and class
+    A group always has a name, a type, a list of children, and a list of dictionaries as attributes.
+    """
+    return dict(name=name, type='group',
+                attributes=[dict(name='NX_class', dtype='string', values=nx_class)],
+                children=[])
+
+
+def _get_or_add_group(children: list, name: str, nx_class: str):
+    """Get or add a group with the specified name and class to a list of children"""
+    g = [x for x in children if _is_group(x, name)]
+    if len(g):
+        return g[0], children
+    children.append(_make_group(name, nx_class))
+    return children[-1], children
+
+
+def _get_or_add_stream(children: list, name: str, stream_config: dict):
+    """Get or add a stream with the specified name and config to a list of children"""
+    m = [x for x in children if _is_stream(x, name)]
+    if len(m):
+        # check that the stream-config is right?
+        return m[0], children
+    children.append(dict(module=name, config=stream_config))
+    return children[-1], children
+
+
+# def a_log(ch: dict):
+#     """Unused, temporarily kept for debugging. May have been correct at some point, wrong now."""
+#     attrs = dict(name='NX_class', type='string', values='NXlog')
+#     units = dict(name='units', type='string', values=ch.get('units', 'dimensionless'))
+#     log_child = dict(module='f144', source=ch['source'], topic=ch['topic'], type=ch['dtype'], attributes=[units])
+#     # log_child = dict(module='f144', source=ch['source'], topic=ch['topic'], dtype=ch['dtype'], attributes=[units])
+#     desc_child = dict(module='dataset', config=dict(name='description', values=ch['description'], type='string'))
+#     return dict(name=ch['name'], type='group', attributes=[attrs], children=[log_child, desc_child])
+
+
+def a_log_as_of_20230626(ch: dict):
+    """Correct form as of June 26, 2023. Notably, source, topic, type, and unit go in a config field.
+
+    The ch dict must have the following keys:
+        - name: the name of the logged value
+        - dtype: the data type of the logged value
+        - source: the Kafka source of the logged value
+        - topic: the Kafka topic of the logged value
+        - description: a description of the logged value
+        - module: the flatbuffer module to use to log the value, e.g., 'f144'
+        - unit: the unit of the logged value, e.g., 'dimensionless'
+
+    The returned structure is:
+        {name: <name>, type: 'group', attributes: [{name: 'NX_class', type: 'string', values: 'NXlog'}],
+         children: [
+          {module: <module>, config: {type: <dtype>, source: <source>, topic: <topic>, unit: <unit>}},
+          {module: 'dataset', config: {name: 'description', type: 'string', values: <description>}}
+         ]
+        }
+    """
+    c = dict(type=ch['dtype'], topic=ch['topic'], source=ch['source'], unit=ch.get('unit', 'dimensionless'))
+    # Use f144 for most things, or f143 for more general objects -- like strings
+    log_child = dict(module=ch['module'], config=c)
+    attrs = dict(name='NX_class', type='string', values='NXlog')
+    desc_child = dict(module='dataset', config=dict(name='description', values=ch['description'], type='string'))
+    return dict(name=ch['name'], type='group', attributes=[attrs], children=[log_child, desc_child])
+
+
+def default_nexus_structure(instr, origin: str | None = None):
+    from zenlog import log
+    import moreniius.additions  # patches the Instance class to have more translation methods
+    from moreniius import MorEniius
+    log.info('Creating NeXus structure from instrument'
+             ' -- no custom Instance to NeXus mapping is used'
+             ' -- provide a JSON object, a python module and function name, or executable to use a custom mapping')
+    return MorEniius.from_mccode(instr, origin=origin, only_nx=False, absolute_depends_on=True).to_nexus_structure()
+
+
+def add_pvs_to_nexus_structure(ns: dict, pvs: list[dict]):
+    if 'children' not in ns:
+        raise RuntimeError('Top-level NeXus structure dict with toplevel list entry named "children" expected.')
+    entry, ns['children'] = _get_or_add_group(ns['children'], 'entry', 'NXentry')
+    # # NXlogs isn't a NeXus base class ...
+    # logs, entry['children'] = _get_or_add_group(entry['children'],  'logs', 'NXlogs')
+    # So dump everything directly into 'children'
+    # but 'NXparameters' _does_ exist:
+    parameters, entry['children'] = _get_or_add_group(entry['children'], 'parameters', 'NXparameters')
+    for pv in pvs:
+        if any(x not in pv for x in ['name', 'dtype', 'source', 'topic', 'description', 'module', 'unit']):
+            raise RuntimeError(f"PV {pv['name']} is missing one or more required keys")
+        parameters['children'].append(a_log_as_of_20230626(pv))
+    return ns
+
+
+def add_title_to_nexus_structure(ns: dict, title: str):
+    if 'children' not in ns:
+        raise RuntimeError('Top-level NeXus structure dict with toplevel list entry named "children" expected.')
+    entry, ns['children'] = _get_or_add_group(ns['children'], 'entry', 'NXentry')
+    entry['children'].append(dict(module='dataset', config=dict(name='title', values=title, type='string')))
+    return ns
+
+
+def insert_events_in_nexus_structure(ns: dict, config: dict):
+    if 'children' not in ns:
+        raise RuntimeError('Top-level NeXus structure dict with toplevel list entry named "children" expected.')
+    entry, ns['children'] = _get_or_add_group(ns['children'], 'entry', 'NXentry')
+
+    # check whether 'instrument' is already a group under 'entry', and add it if not
+    instr, entry['children'] = _get_or_add_group(entry['children'], 'instrument', 'NXinstrument')
+
+    # check whether 'detector' is a group under '/entry/instrument', and add it if not
+    detector, instr['children'] = _get_or_add_group(instr['children'], 'detector', 'NXdetector')
+    # ... TODO fill in all of the required detector elements :(
+
+    # check whether 'events' is a group under `/entry/instrument/detector`
+    events, detector['children'] = _get_or_add_group(detector['children'], 'events', 'NXevent_data')
+
+    # Ensure that the events group has the correct stream-specification child
+    # {'module': 'ev44', 'config': {'source': 'source', 'topic': 'topic'}}
+    stream, events['children'] = _get_or_add_stream(events['children'], 'ev44', config)
+
+    return ns
+
+
+def get_writer_pool(broker: str = None, job: str = None, command: str = None):
+    from file_writer_control import WorkerJobPool
+    pool = WorkerJobPool(f"{broker}/{job}", f"{broker}/{command}")
+    return pool
+
+
+def define_nexus_structure(instr: Union[Path, str], pvs: list[dict], title: str = None, event_stream: dict[str, str] = None,
+                           file: Union[Path, None] = None, func: Union[Callable[[Instr], dict], None] = None,
+                           binary: Union[Path, None] = None, origin: str = None):
+    import json
+    from .mccode import get_mcstas_instr
+    if file is not None and file.exists():
+        with open(file, 'r') as file:
+            nexus_structure = json.load(file)
+    elif func is not None:
+        nexus_structure = func(get_mcstas_instr(instr))
+    elif binary is not None and binary.exists():
+        from subprocess import run, PIPE
+        result = run([binary, str(instr)], stdout=PIPE, stderr=PIPE)
+        if result.returncode != 0:
+            raise RuntimeError(f"Failed to execute {binary} {instr} due to error {result.stderr.decode()}")
+        nexus_structure = json.loads(result.stdout.decode())
+    else:
+        nexus_structure = default_nexus_structure(get_mcstas_instr(instr), origin=origin)
+    nexus_structure = add_pvs_to_nexus_structure(nexus_structure, pvs)
+    nexus_structure = add_title_to_nexus_structure(nexus_structure, title)
+    # nexus_structure = insert_events_in_nexus_structure(nexus_structure, event_stream)
+    return nexus_structure
+
+
+def start_pool_writer(start_time_string, structure, filename=None, stop_time_string: str | None = None,
+                      broker: str | None = None, job_topic: str | None = None, command_topic: str | None = None,
+                      wait: bool = False, timeout: float | None = None, job_id: str | None = None):
+    from sys import exit
+    from os import EX_OK, EX_UNAVAILABLE
+    from time import sleep
+    from json import dumps
+    from datetime import datetime, timedelta
+    from file_writer_control import JobHandler, WriteJob, CommandState
+
+    start_time = datetime.fromisoformat(start_time_string)
+    if filename is None:
+        filename = f'{start_time:%Y%m%d_%H%M%S}.nxs'
+
+    pool = get_writer_pool(broker=broker, job=job_topic, command=command_topic)
+    handler_opts = {'worker_finder': pool}
+
+    handler = JobHandler(**handler_opts)
+    small_string = dumps(structure, indent=None, separators=(',', ':'))
+
+    end_time = datetime.now() if wait else None
+    if stop_time_string is not None:
+        end_time = datetime.fromisoformat(stop_time_string)
+    print(f"write file from {start_time} until {end_time}")
+
+    job = WriteJob(small_string, filename, broker, start_time, end_time, job_id=job_id or "")
+    # start the job
+    start = handler.start_job(job)
+    if timeout is not None:
+        try:
+            # ensure the start succeeds:
+            zero_time = datetime.now()
+            while not start.is_done():
+                if zero_time + timedelta(seconds=timeout) < datetime.now():
+                    raise RuntimeError(f"Timed out while starting job {job.job_id}")
+                elif start.get_state() == CommandState.ERROR:
+                    raise RuntimeError(f"Starting job {job.job_id} failed with message {start.get_message()}")
+                sleep(1)
+        except RuntimeError as e:
+            # raise RuntimeError(e.__str__() + f" The message was: {start.get_message()}")
+            print(f"{e} The message was: {start.get_message()}")
+            exit(EX_UNAVAILABLE)
+
+    if wait:
+        try:
+            while not handler.is_done():
+                sleep(1)
+        except RuntimeError as error:
+            print(str(error) + f'Writer failed, producing message:\n{handler.get_message}')
+            exit(EX_UNAVAILABLE)
+    exit(EX_OK)
+
+
+def get_arg_parser():
+    from argparse import ArgumentParser
+    from .utils import is_callable, is_readable, is_executable, is_writable
+    parser = ArgumentParser(description="Control writing Kafka stream(s) to a NeXus file")
+    a = parser.add_argument
+    a('instrument', type=str, default=None, help="The mcstas instrument with EPICS PVs")
+    a('-p', '--prefix', type=str, default='mcstas:')
+    a('-t', '--topic', type=str, help="The Kafka broker topic to instruct the Forwarder to use")
+    a('-b', '--broker', type=str, help="The Kafka broker server used by the Writer")
+    a('-j', '--job', type=str, help='Writer job topic')
+    a('-c', '--command', type=str, help='Writer command topic')
+    a('--title', type=str, default='scan title for testing', help='Output file title parameter')
+    a('--event-source', type=str)
+    a('--event-topic', type=str)
+    a('-f', '--filename', type=str, default=None)
+    a('--ns-func', type=is_callable, default=None, help='Python module:function to produce NeXus structure')
+    a('--ns-file', type=is_readable, default=None, help='Base NeXus structure, will be extended')
+    a('--ns-exec', type=is_executable, default=None, help='Executable to produce NeXus structure')
+    a('--ns-save', type=is_writable, default=None, help='Path at which to save (overwrite) extended NeXus structure')
+    a('--start-time', type=str)
+    a('--stop-time',  type=str, default=None)
+    a('--origin', type=str, default=None, help='component name used for the origin of the NeXus file')
+    a('--wait', action='store_true', help='If provided, wait for the writer to finish before exiting')
+    a('--time-out', type=float, default=120., help='Wait up to the timeout for writing to start')
+    a('--job-id', type=str, default=None, help='Unique Job identifier for this write-job')
+
+    return parser
+
+
+def parameter_description(inst_param):
+    desc = f"{inst_param.value.data_type} valued McStas parameter '{inst_param.name}', "
+    desc += f"default: {inst_param.value}" if inst_param.value.has_value else "no default"
+    if inst_param.unit is not None:
+        desc += f" and expected units of {inst_param.unit}"
+    return desc
+
+
+def construct_writer_pv_dicts(instr: Union[Path, str], prefix: str, topic: str):
+    from .mccode import get_mccode_instr_parameters
+    parameters = get_mccode_instr_parameters(instr)
+    return construct_writer_pv_dicts_from_parameters(parameters, prefix, topic)
+
+
+def construct_writer_pv_dicts_from_parameters(parameters, prefix: str, topic: str):
+    def strip_quotes(s):
+        return s[1:-1] if s is not None and len(s) > 2 and (s[0] == s[-1] == '"' or s[0] == s[-1] == "'") else s
+    return [dict(name=p.name, dtype=p.value.data_type.name, source=f'{prefix}{p.name}', topic=topic,
+                 description=parameter_description(p), module='f144', unit=strip_quotes(p.unit)) for p in parameters]
+
+
+def parse_writer_args():
+    args = get_arg_parser().parse_args()
+    params = construct_writer_pv_dicts(args.instrument, args.prefix, args.topic)
+    structure = define_nexus_structure(args.instrument, params, title=args.title, origin=args.origin,
+                                       file=args.ns_file, func=args.ns_func, binary=args.ns_exec,
+                                       event_stream={'source': args.event_source, 'topic': args.event_topic})
+    if args.ns_save is not None:
+        from json import dump
+        with open(args.ns_save, 'w') as file:
+            dump(structure, file, indent=2)
+
+    return args, params, structure
+
+
+def print_time():
+    from datetime import datetime
+    print(datetime.now())
+
+
+def start_writer():
+    args, parameters, structure = parse_writer_args()
+    return start_pool_writer(args.start_time, structure, args.filename, stop_time_string=args.stop_time,
+                             broker=args.broker, job_topic=args.job, command_topic=args.command,
+                             wait=args.wait, timeout=args.time_out, job_id=args.job_id)
+
+
+def wait_on_writer():
+    from sys import exit
+    from os import EX_OK, EX_UNAVAILABLE
+    from time import sleep
+    from datetime import datetime, timedelta
+    from file_writer_control import JobHandler, CommandState
+
+    from argparse import ArgumentParser
+    parser = ArgumentParser()
+    a = parser.add_argument
+    a('-b', '--broker', type=str, help="The Kafka broker server used by the Writer")
+    a('-j', '--job', type=str, help='Writer job topic')
+    a('-c', '--command', type=str, help='Writer command topic')
+    a('id', type=str, help='Job id to wait on')
+    a('-s', '--stop-after', type=float, help='Stop after time, seconds', default=1)
+    a('-t', '--time-out', type=float, help='Time out after, seconds', default=24*60*60*30)
+    args = parser.parse_args()
+
+    pool = get_writer_pool(broker=args.broker, job=args.job, command=args.command)
+    job = JobHandler(worker_finder=pool, job_id=args.id)
+    stop_time = datetime.now() + timedelta(seconds=args.stop_after)
+    stop = job.set_stop_time(stop_time)
+
+    try:
+        timeout = args.time_out
+        zero_time = datetime.now()
+        while not stop.is_done() and not job.is_done():
+            if zero_time + timedelta(seconds=timeout) < datetime.now():
+                print('1')
+                raise RuntimeError(f"Timed out while stopping job {job.job_id}")
+            elif stop.get_state() == CommandState.ERROR:
+                print('2')
+                raise RuntimeError(f"Stopping job {job.job_id} failed with message {stop.get_message()}")
+            sleep(0.5)
+    except RuntimeError as e:
+        # raise RuntimeError(e.__str__() + f" The message was: {stop.get_message()}")
+        exit(EX_UNAVAILABLE)
+    exit(EX_OK)

mccode_plumber-0.6.0/src/mccode_plumber.egg-info/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: mccode-plumber
+Version: 0.6.0
+Author-email: Gregory Tucker <gregory.tucker@ess.eu>
+Classifier: License :: OSI Approved :: BSD License
+Description-Content-Type: text/markdown
+Requires-Dist: p4p
+Requires-Dist: file-writer-control>=1.3.0
+Requires-Dist: restage>=0.4.0
+Requires-Dist: mccode-to-kafka>=0.2.1
+Requires-Dist: moreniius>=0.2.3
+Requires-Dist: icecream
+
+# McCode Plumber
+Setup, run, and teardown the infrastructure for splitrun McCode scans sending data through Kafka into NeXus

mccode_plumber-0.6.0/src/mccode_plumber.egg-info/SOURCES.txt

@@ -0,0 +1,23 @@
+.gitignore
+README.md
+pyproject.toml
+.github/workflows/pip.yml
+.github/workflows/wheels.yml
+src/mccode_plumber/__init__.py
+src/mccode_plumber/conductor.py
+src/mccode_plumber/epics.py
+src/mccode_plumber/forwarder.py
+src/mccode_plumber/kafka.py
+src/mccode_plumber/mccode.py
+src/mccode_plumber/splitrun.py
+src/mccode_plumber/utils.py
+src/mccode_plumber/writer.py
+src/mccode_plumber.egg-info/PKG-INFO
+src/mccode_plumber.egg-info/SOURCES.txt
+src/mccode_plumber.egg-info/dependency_links.txt
+src/mccode_plumber.egg-info/entry_points.txt
+src/mccode_plumber.egg-info/requires.txt
+src/mccode_plumber.egg-info/top_level.txt
+tests/test_epics.py
+tests/test_splitrun.py
+tests/test_writer.py

mccode_plumber-0.6.0/src/mccode_plumber.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

mccode_plumber-0.6.0/src/mccode_plumber.egg-info/entry_points.txt

@@ -0,0 +1,11 @@
+[console_scripts]
+mp-epics = mccode_plumber.epics:run
+mp-epics-update = mccode_plumber.epics:update
+mp-forwarder-setup = mccode_plumber.forwarder:setup
+mp-forwarder-teardown = mccode_plumber.forwarder:teardown
+mp-insert-hdf5-instr = mccode_plumber.mccode:insert
+mp-register-topics = mccode_plumber.kafka:register_topics
+mp-splitrun = mccode_plumber.splitrun:main
+mp-writer-from = mccode_plumber.writer:print_time
+mp-writer-wait = mccode_plumber.writer:wait_on_writer
+mp-writer-write = mccode_plumber.writer:start_writer

mccode_plumber-0.6.0/src/mccode_plumber.egg-info/requires.txt

@@ -0,0 +1,6 @@
+p4p
+file-writer-control>=1.3.0
+restage>=0.4.0
+mccode-to-kafka>=0.2.1
+moreniius>=0.2.3
+icecream

mccode_plumber-0.6.0/src/mccode_plumber.egg-info/top_level.txt

@@ -0,0 +1 @@
+mccode_plumber

mccode_plumber-0.6.0/tests/test_epics.py

@@ -0,0 +1,48 @@
+import unittest
+
+
+class EPICSTestCase(unittest.TestCase):
+    def setUp(self):
+        from uuid import uuid4
+        from multiprocessing import Process
+        from mccode_plumber.epics import main, convert_instr_parameters_to_nt
+        from mccode_antlr.loader.loader import parse_mccode_instr_parameters, parse_mcstas_instr
+        instr = 'define instrument blah(par1, double par2, int par3=1, string par4="string", double par5=5.5) trace end'
+        self.pars = parse_mccode_instr_parameters(instr)
+        self.pvs = convert_instr_parameters_to_nt(parse_mcstas_instr(instr).parameters)
+        self.prefix = f"test{str(uuid4()).replace('-', '')}:"
+        self.proc = Process(target=main, args=(self.pvs, self.prefix))
+        self.proc.start()
+
+    def tearDown(self):
+        self.proc.terminate()
+        self.proc.join(1)
+        self.proc.close()
+
+    def test_server_runs(self):
+        from p4p.client.thread import Context
+        providers = Context.providers()
+        self.assertTrue('pva' in providers)
+        ctx = Context('pva')
+
+        for par in self.pars:
+            pv = ctx.get(f"{self.prefix}{par.name}")
+            self.assertTrue(pv is not None)
+            if par.value.has_value:
+                self.assertEqual(pv, par.value.value)
+
+    def test_update_pvs(self):
+        from p4p.client.thread import Context
+        ctx = Context('pva')
+        values = {'par1': 1.1, 'par2': 2.2, 'par3': 3, 'par4': 'four', 'par5': 55.555}
+        for name, value in values.items():
+            ctx.put(f"{self.prefix}{name}", value)
+
+        for name, value in values.items():
+            pv = ctx.get(f"{self.prefix}{name}")
+            self.assertTrue(pv is not None)
+            self.assertEqual(pv, value)
+
+
+if __name__ == '__main__':
+    unittest.main()

mccode_plumber-0.6.0/tests/test_splitrun.py

@@ -0,0 +1,46 @@
+import unittest
+
+
+class SplitrunTestCase(unittest.TestCase):
+    def test_parsing(self):
+        from mccode_plumber.splitrun import make_parser
+        parser = make_parser()
+        args = parser.parse_args(['--broker', 'l:9092', '--source', 'm', '-n', '10000', 'inst.h5', '--', 'a=1:4', 'b=2:5'])
+        self.assertEqual(args.instrument, ['inst.h5'])
+        self.assertEqual(args.broker, 'l:9092')
+        self.assertEqual(args.source, 'm')
+        self.assertEqual(args.ncount, [10000])
+        self.assertEqual(args.parameters, ['a=1:4', 'b=2:5'])
+        self.assertFalse(args.parallel)
+
+    def test_mixed_order_throws(self):
+        from mccode_plumber.splitrun import make_parser
+        parser = make_parser()
+        parser.prog = "{{This is supposed to fail, don't be dismayed by this output!}}"
+        # These also output usage information to stdout -- don't be surprised by the 'extra' test output.
+        with self.assertRaises(SystemExit):
+            parser.parse_args(['inst.h5', '--broker', 'l:9092', '--source', 'm', '-n', '10000', 'a=1:4', 'b=2:5'])
+        with self.assertRaises(SystemExit):
+            parser.parse_args(['--broker', 'l:9092', '--source', 'm', 'inst.h5', '-n', '10000', 'a=1:4', 'b=2:5'])
+
+    def test_sort_args(self):
+        from mccode_antlr.run.runner import sort_args
+        self.assertEqual(sort_args(['-n', '10000', 'inst.h5', 'a=1:4', 'b=2:5']), ['-n', '10000', 'inst.h5', 'a=1:4', 'b=2:5'])
+        self.assertEqual(sort_args(['inst.h5', '-n', '10000', 'a=1:4', 'b=2:5']), ['-n', '10000', 'inst.h5', 'a=1:4', 'b=2:5'])
+
+    def test_sorted_mixed_order_does_not_throw(self):
+        from mccode_plumber.splitrun import make_parser
+        from mccode_antlr.run.runner import sort_args
+        parser = make_parser()
+        args = parser.parse_args(sort_args(['inst.h5', '--broker', 'www.github.com:9093', '--source', 'dev/null',
+                                            '-n', '123', '--parallel', '--', 'a=1:4', 'b=2:5']))
+        self.assertEqual(args.instrument, ['inst.h5'])
+        self.assertEqual(args.broker, 'www.github.com:9093')
+        self.assertEqual(args.source, 'dev/null')
+        self.assertEqual(args.ncount, [123])
+        self.assertEqual(args.parameters, ['a=1:4', 'b=2:5'])
+        self.assertTrue(args.parallel)
+
+
+if __name__ == '__main__':
+    unittest.main()

mccode_plumber-0.6.0/tests/test_writer.py

@@ -0,0 +1,71 @@
+import unittest
+
+
+class WriterTestCase(unittest.TestCase):
+    def setUp(self):
+        from json import dumps
+        from mccode_antlr.loader import parse_mcstas_instr
+        from mccode_to_kafka.writer import da00_dataarray_config, da00_variable_config
+        t = da00_variable_config(name='t', label='monitor', unit='usec', axes=['t'], shape=[10], data_type='float64')
+        ns = da00_dataarray_config(topic='monitor', source='mccode-to-kafka', variables=[t])
+        instr = f"""DEFINE INSTRUMENT this_IS_NOT_BIFROST()
+        TRACE
+        COMPONENT origin = Arm() AT (0, 0, 0) ABSOLUTE
+        COMPONENT source = Source_simple() AT (0, 0, 1) RELATIVE PREVIOUS
+        COMPONENT monitor = TOF_monitor() AT (0, 0, 1) RELATIVE source
+        METADATA "application/json" "nexus_structure_stream_data" %{{{dumps(ns)}%}}
+        COMPONENT sample = Arm() AT (0, 0, 80) RELATIVE source
+        END
+        """
+        self.instr = parse_mcstas_instr(instr)
+
+    def test_parse(self):
+        from mccode_plumber.writer import construct_writer_pv_dicts_from_parameters
+        from mccode_plumber.writer import default_nexus_structure
+        params = construct_writer_pv_dicts_from_parameters(self.instr.parameters, 'mcstas:', 'topic')
+        self.assertEqual(len(params), 0)
+        struct = default_nexus_structure(self.instr)
+
+        self.assertEqual(len(struct['children']), 1)
+        self.assertEqual(struct['children'][0]['name'], 'entry')
+        self.assertEqual(struct['children'][0]['children'][0]['name'], 'instrument')
+        self.assertEqual(struct['children'][0]['children'][0]['children'][1]['name'], '0_origin')
+        self.assertEqual(struct['children'][0]['children'][0]['children'][2]['name'], '1_source')
+        self.assertEqual(struct['children'][0]['children'][0]['children'][3]['name'], '2_monitor')
+        mon = struct['children'][0]['children'][0]['children'][3]
+        self.assertEqual(len(mon['children']), 4)  # removed 'mccode' property 5->4
+        idx = [i for i, ch in enumerate(mon['children']) if 'name' in ch and 'data' == ch['name']]
+        self.assertTrue(len(idx), 1)
+        data = mon['children'][idx[0]]
+        idx = [i for i, ch in enumerate(data['children']) if 'module' in ch and 'da00' == ch['module']]
+        self.assertEqual(len(idx), 1)
+        da00 = data['children'][idx[0]]
+        self.assertEqual(len(da00.keys()), 2)
+        self.assertEqual(da00['module'], 'da00')
+        self.assertEqual(da00['config']['topic'], 'monitor')
+        self.assertEqual(da00['config']['source'], 'mccode-to-kafka')
+
+
+class WriterUnitsTestCase(unittest.TestCase):
+    def setUp(self):
+        from mccode_antlr.loader import parse_mcstas_instr
+        instr = f"""DEFINE INSTRUMENT with_logs(double a/"Hz", b/"m", int c, string d)
+        TRACE
+        COMPONENT origin = Arm() AT (0, 0, 0) ABSOLUTE
+        COMPONENT source = Source_simple() AT (0, 0, 1) RELATIVE PREVIOUS
+        COMPONENT sample = Arm() AT (0, 0, 80) RELATIVE source
+        END
+        """
+        self.instr = parse_mcstas_instr(instr)
+
+    def test_parse(self):
+        from mccode_plumber.writer import construct_writer_pv_dicts_from_parameters
+        params = construct_writer_pv_dicts_from_parameters(self.instr.parameters, 'mcstas:', 'topic')
+        self.assertEqual(len(params), 4)
+        for p, x in zip(params, [('a', 'Hz'), ('b', 'm'), ('c', None), ('d', None)]):
+            self.assertEqual(p['name'], x[0])
+            self.assertEqual(p['unit'], x[1])
+
+
+if __name__ == '__main__':
+    unittest.main()