From 9c6c0133a627289587f907ccdd2ffbb058751c29 Mon Sep 17 00:00:00 2001 From: MatInGit Date: Tue, 7 Apr 2026 11:37:58 +0100 Subject: [PATCH 1/4] intial provenance pass --- .../local/deployment_config_event_driven.yaml | 86 ++ ...loyment_config_multihead_event_driven.yaml | 130 +++ .../local/model_definition_event_driven.py | 51 + poly_lithic/src/cli.py | 47 +- poly_lithic/src/config/config_object.py | 11 + poly_lithic/src/interfaces/BaseInterface.py | 2 + .../src/interfaces/fastapi_interface.py | 34 +- poly_lithic/src/interfaces/k2eg_interface.py | 1 + poly_lithic/src/interfaces/p4p_interface.py | 102 +- poly_lithic/src/utils/builder.py | 28 +- poly_lithic/src/utils/messaging.py | 131 ++- poly_lithic/src/utils/trace_server.py | 69 ++ poly_lithic/src/utils/trace_store.py | 107 ++ poly_lithic/src/utils/trace_viewer.html | 450 +++++++++ tests/test_event_driven.py | 917 ++++++++++++++++++ tests/test_interfaceobserver.py | 4 +- 16 files changed, 2118 insertions(+), 52 deletions(-) create mode 100644 examples/base/local/deployment_config_event_driven.yaml create mode 100644 examples/base/local/deployment_config_multihead_event_driven.yaml create mode 100644 examples/base/local/model_definition_event_driven.py create mode 100644 poly_lithic/src/utils/trace_server.py create mode 100644 poly_lithic/src/utils/trace_store.py create mode 100644 poly_lithic/src/utils/trace_viewer.html create mode 100644 tests/test_event_driven.py diff --git a/examples/base/local/deployment_config_event_driven.yaml b/examples/base/local/deployment_config_event_driven.yaml new file mode 100644 index 0000000..440b8bc --- /dev/null +++ b/examples/base/local/deployment_config_event_driven.yaml @@ -0,0 +1,86 @@ +# Event-driven deployment example +# Instead of polling at a fixed rate, the pipeline triggers automatically +# whenever a monitored PV value changes. +# +# Usage: +# pl run --config ./examples/base/local/deployment_config_event_driven.yaml +# +# Test by writing to a PV: +# pvput ML:LOCAL:TEST_A 3.14 +# pvput ML:LOCAL:TEST_B 2.71 +# +# The model will only evaluate when an input PV changes, rather than +# every N seconds. + +deployment: + type: "event_driven" + min_monitor_interval: 0.01 # throttle: at most one update per PV every 100 ms + on_change_only: true # skip updates where the value hasn't changed + +modules: + p4p_server: + name: "p4p_server" + type: "interface.p4p_server" + pub: + - "in_interface" + sub: + - "get_all" + - "out_transformer" + module_args: None + config: + EPICS_PVA_NAME_SERVERS: "localhost:5075" + variables: + ML:LOCAL:TEST_B: + proto: pva + name: ML:LOCAL:TEST_B + ML:LOCAL:TEST_A: + proto: pva + name: ML:LOCAL:TEST_A + ML:LOCAL:TEST_S: + proto: pva + name: ML:LOCAL:TEST_S + + input_transformer: + name: "input_transformer" + type: "transformer.SimpleTransformer" + pub: "in_transformer" + sub: "in_interface" + module_args: None + config: + symbols: + - "ML:LOCAL:TEST_B" + - "ML:LOCAL:TEST_A" + variables: + x: + formula: "ML:LOCAL:TEST_A * 2 + 10" + y: + formula: "ML:LOCAL:TEST_B + 120" + + model: + name: "model" + type: "model.SimpleModel" + pub: "model" + sub: "in_transformer" + module_args: None + config: + type: "LocalModelGetter" + args: + model_path: "examples/base/local/model_definition_event_driven.py" + model_factory_class: "ModelFactory" + variables: + max: + type: "scalar" + + output_transformer: + name: "output_transformer" + type: "transformer.SimpleTransformer" + pub: "out_transformer" + sub: "model" + module_args: + unpack_data: True + config: + symbols: + - "output" + variables: + ML:LOCAL:TEST_S: + formula: "output" diff --git a/examples/base/local/deployment_config_multihead_event_driven.yaml b/examples/base/local/deployment_config_multihead_event_driven.yaml new file mode 100644 index 0000000..772a148 --- /dev/null +++ b/examples/base/local/deployment_config_multihead_event_driven.yaml @@ -0,0 +1,130 @@ +# Multi-headed event-driven deployment example +# Two input interfaces (FastAPI + p4p server) feed the same pipeline. +# The model evaluates whenever EITHER interface receives new values. +# +# Usage: +# pl run --config ./examples/base/local/deployment_config_multihead_event_driven.yaml -p +# +# Test via PVAccess: +# spput ML:LOCAL:TEST_A 3.14 +# spput ML:LOCAL:TEST_B 2.71 +# +# Test via HTTP: +# curl -X POST http://localhost:8001/submit -H 'Content-Type: application/json' -d '{"variables":{"ML:LOCAL:TEST_A":{"value":3.14},"ML:LOCAL:TEST_B":{"value":2.71}}}' +# +# Monitor output PV: +# spmonitor ML:LOCAL:TEST_S + +deployment: + type: "event_driven" + min_monitor_interval: 0.01 # throttle per PV (seconds) + on_change_only: false + +modules: + # ── PVAccess server interface ──────────────────────────────────────── + p4p_server: + name: "p4p_server" + type: "interface.p4p_server" + pub: + - "in_interface_1" + sub: + - "get_all" + - "out_transformer" + - "in_interface_0" + module_args: None + config: + EPICS_PVA_NAME_SERVERS: "localhost:5075" + variables: + ML:LOCAL:TEST_A: + proto: pva + name: ML:LOCAL:TEST_A + ML:LOCAL:TEST_B: + proto: pva + name: ML:LOCAL:TEST_B + ML:LOCAL:TEST_S: + proto: pva + name: ML:LOCAL:TEST_S + + # ── FastAPI HTTP interface ─────────────────────────────────────────── + fastapi_server: + name: "fastapi_server" + type: "interface.fastapi_server" + pub: + - "in_interface_0" + sub: + - "get_all" + - "out_transformer" + - "in_interface_1" + module_args: None + config: + name: "fastapi_server" + host: "0.0.0.0" + port: 8001 + start_server: true + wait_for_server_start: true + startup_timeout_s: 5.0 + input_queue_max: 1000 + output_queue_max: 1000 + variables: + ML:LOCAL:TEST_A: + mode: in + type: scalar + default: 0.0 + ML:LOCAL:TEST_B: + mode: in + type: scalar + default: 0.0 + ML:LOCAL:TEST_S: + mode: out + type: scalar + default: 0.0 + + # ── Input transformer ─────────────────────────────────────────────── + input_transformer: + name: "input_transformer" + type: "transformer.SimpleTransformer" + pub: "in_transformer" + sub: + - "in_interface_0" + - "in_interface_1" + module_args: None + config: + symbols: + - "ML:LOCAL:TEST_A" + - "ML:LOCAL:TEST_B" + variables: + x: + formula: "ML:LOCAL:TEST_A" + y: + formula: "ML:LOCAL:TEST_B" + + # ── Model ──────────────────────────────────────────────────────────── + model: + name: "model" + type: "model.SimpleModel" + pub: "model" + sub: "in_transformer" + module_args: None + config: + type: "LocalModelGetter" + args: + model_path: "examples/base/local/model_definition_event_driven.py" + model_factory_class: "ModelFactory" + variables: + max: + type: "scalar" + + # ── Output transformer ────────────────────────────────────────────── + output_transformer: + name: "output_transformer" + type: "transformer.SimpleTransformer" + pub: "out_transformer" + sub: "model" + module_args: + unpack_data: True + config: + symbols: + - "output" + variables: + ML:LOCAL:TEST_S: + formula: "output" diff --git a/examples/base/local/model_definition_event_driven.py b/examples/base/local/model_definition_event_driven.py new file mode 100644 index 0000000..d463ad6 --- /dev/null +++ b/examples/base/local/model_definition_event_driven.py @@ -0,0 +1,51 @@ +import torch +import os +import time +import logging + +logger = logging.getLogger(__name__) + + +class ModelFactory: + def __init__(self): + os.environ['PYTHONPATH'] = os.path.abspath( + os.path.join(os.path.dirname(__file__), '..', '..', '..') + ) + self.model = SimpleModel() + model_path = 'examples/base/local/model.pth' + if os.path.exists(model_path): + self.model.load_state_dict(torch.load(model_path)) + logger.info('Model loaded successfully.') + else: + logger.warning( + f"Model file '{model_path}' not found. Using untrained model." + ) + logger.info('ModelFactory initialized (event-driven mode)') + + def get_model(self): + return self.model + + +class SimpleModel(torch.nn.Module): + def __init__(self): + super(SimpleModel, self).__init__() + self.linear1 = torch.nn.Linear(2, 10) + self.linear2 = torch.nn.Linear(10, 1) + self._eval_count = 0 + + def forward(self, x): + x = torch.relu(self.linear1(x)) + x = self.linear2(x) + return x + + def evaluate(self, x: dict) -> dict: + self._eval_count += 1 + logger.info( + f'[event-driven] evaluate #{self._eval_count} triggered at ' + f'{time.strftime("%H:%M:%S")} | inputs: x={x.get("x")}, y={x.get("y")}' + ) + input_tensor = torch.tensor([x['x'], x['y']], dtype=torch.float32) + output_tensor = self.forward(input_tensor) + result = output_tensor.item() + logger.info(f'[event-driven] output: {result}') + return {'output': result} diff --git a/poly_lithic/src/cli.py b/poly_lithic/src/cli.py index 4461777..f61d709 100644 --- a/poly_lithic/src/cli.py +++ b/poly_lithic/src/cli.py @@ -71,7 +71,7 @@ def load_env_config(env_path): raise e -async def model_main(args, config, broker): +async def model_main(args, config, broker, builder=None): """ Main async function for running the model manager. @@ -79,6 +79,7 @@ async def model_main(args, config, broker): args: Parsed arguments namespace config: Configuration object broker: Broker instance + builder: Builder instance (needed for event_driven mode) """ logger = get_logger() logger.info('Starting model manager') @@ -106,6 +107,37 @@ async def model_main(args, config, broker): await asyncio.sleep(0.01) + elif config.deployment.type == 'event_driven': + # Start PV monitors on all input interfaces + input_observers = builder.get_input_interface_observers() + if not input_observers: + raise ValueError('No input InterfaceObservers found for event-driven mode') + + # Seed all transformer inputs with current PV values so the + # all-inputs-present gate passes on the first monitor event. + broker.get_all() + while broker.queue: + broker.parse_queue() + logger.info('Seeded transformer inputs via get_all') + + for obs in input_observers: + obs.start_monitors( + broker, + min_interval=config.deployment.min_monitor_interval, + on_change_only=config.deployment.on_change_only, + ) + logger.info(f'Started monitors on {len(input_observers)} input interface(s)') + + while True: + if len(broker.queue) > 0: + broker.parse_queue() + + if args.one_shot: + logger.info('One shot mode, exiting') + break + + await asyncio.sleep(0.01) + else: raise Exception(f'Deployment type "{config.deployment.type}" not supported') @@ -218,10 +250,19 @@ def run_model(config, model_getter, debug, env, one_shot, publish, requirements) # Import heavy dependencies only when needed from poly_lithic.src.utils.builder import Builder + from poly_lithic.src.utils.trace_store import TraceStore + from poly_lithic.src.utils.trace_server import start_trace_server click.echo('Building model manager...') builder = Builder(config) - broker = builder.build() + + trace_store = TraceStore(maxlen=builder.config.deployment.trace_buffer_size) + broker = builder.build(trace_store=trace_store) + + # Start tracing API server + trace_port = int(os.environ.get('TRACE_PORT', builder.config.deployment.trace_port)) + start_trace_server(trace_store, port=trace_port) + logger.info(f'Tracing API server started on port {trace_port}') if requirements: click.echo('Requirements-only mode - exiting after installation') @@ -240,7 +281,7 @@ def run_model(config, model_getter, debug, env, one_shot, publish, requirements) ) logger.info('Starting model manager main loop') - asyncio.run(model_main(args, builder.config, broker)) + asyncio.run(model_main(args, builder.config, broker, builder=builder)) except KeyboardInterrupt: click.echo('\n\nInterrupted by user') diff --git a/poly_lithic/src/config/config_object.py b/poly_lithic/src/config/config_object.py index aba8ffe..43bbdfa 100644 --- a/poly_lithic/src/config/config_object.py +++ b/poly_lithic/src/config/config_object.py @@ -30,6 +30,17 @@ def validate_module_args(cls, v): class DeploymentConfig(pydantic.BaseModel): type: str rate: Optional[Union[float, int]] = None + min_monitor_interval: float = 0.0 + on_change_only: bool = False + trace_buffer_size: int = 10000 + trace_port: int = 8100 + + @pydantic.field_validator('type') + def validate_type(cls, v): + allowed = {'continuous', 'event_driven'} + if v not in allowed: + raise ValueError(f'deployment type must be one of {allowed}, got {v!r}') + return v class ConfigObject(pydantic.BaseModel): diff --git a/poly_lithic/src/interfaces/BaseInterface.py b/poly_lithic/src/interfaces/BaseInterface.py index 54082fb..0e23a8b 100644 --- a/poly_lithic/src/interfaces/BaseInterface.py +++ b/poly_lithic/src/interfaces/BaseInterface.py @@ -18,6 +18,8 @@ def save(self, data, **kwargs): class BaseInterface(ABC): + supports_monitor: bool = False + @abstractmethod def __init__(self, config): pass diff --git a/poly_lithic/src/interfaces/fastapi_interface.py b/poly_lithic/src/interfaces/fastapi_interface.py index cce6031..ae1948e 100644 --- a/poly_lithic/src/interfaces/fastapi_interface.py +++ b/poly_lithic/src/interfaces/fastapi_interface.py @@ -101,6 +101,7 @@ def _numpy_to_native(obj): class SimpleFastAPIInterfaceServer(BaseInterface): + supports_monitor = True """HTTP interface backed by FastAPI with an in-memory job queue. Registered as ``"fastapi_server"`` so it can be referenced in the YAML @@ -141,6 +142,7 @@ def __init__(self, config: dict): # -- monitor callback slot (Stage 2 hook) ----------------------- self._monitor_callback = None + self._monitor_callbacks_per_pv: dict[str, list] = {} # {pv_name: [callbacks]} # -- build FastAPI app ------------------------------------------ self.app = self._build_app() @@ -350,14 +352,24 @@ def _enqueue_jobs(self, jobs_input: list[JobInput]) -> list[dict]: 'updated': updated_vars, }) - # -- Phase 4: fire monitor callback ------------------------- - if self._monitor_callback is not None: - for jid, variables in resolved_jobs: - snapshot = self._jobs[jid]['inputs'] + # -- Phase 4: fire monitor callbacks ------------------------ + for jid, variables in resolved_jobs: + snapshot = self._jobs[jid]['inputs'] + # Global callback (continuous mode) + if self._monitor_callback is not None: try: self._monitor_callback(snapshot) except Exception: logger.exception('Monitor callback error (ignored)') + # Per-PV callbacks (event-driven mode) + for vname, vstruct in snapshot.items(): + for cb in self._monitor_callbacks_per_pv.get(vname, []): + try: + cb(vstruct['value']) + except Exception: + logger.exception( + f'Per-PV monitor callback error for {vname} (ignored)' + ) return accepted @@ -511,9 +523,17 @@ def get_inputs(self) -> list[str]: def get_outputs(self) -> list[str]: return list(self._out_list) - def monitor(self, handler, **kwargs) -> bool: - """Register a single monitor callback. Returns ``True``.""" - self._monitor_callback = handler + def monitor(self, handler, pv_name=None, **kwargs) -> bool: + """Register a monitor callback. + + When *pv_name* is given the callback fires with the scalar value + for that specific variable (event-driven per-PV mode). Without + *pv_name* a single global callback receives the full snapshot. + """ + if pv_name is not None: + self._monitor_callbacks_per_pv.setdefault(pv_name, []).append(handler) + else: + self._monitor_callback = handler return True def close(self): diff --git a/poly_lithic/src/interfaces/k2eg_interface.py b/poly_lithic/src/interfaces/k2eg_interface.py index 4a5ab0e..e04203f 100644 --- a/poly_lithic/src/interfaces/k2eg_interface.py +++ b/poly_lithic/src/interfaces/k2eg_interface.py @@ -17,6 +17,7 @@ class K2EGInterface(BaseInterface): + supports_monitor = True def __init__(self, config): try: self.client = k2eg.dml( diff --git a/poly_lithic/src/interfaces/p4p_interface.py b/poly_lithic/src/interfaces/p4p_interface.py index 9d3b511..135c256 100644 --- a/poly_lithic/src/interfaces/p4p_interface.py +++ b/poly_lithic/src/interfaces/p4p_interface.py @@ -28,6 +28,8 @@ class SimplePVAInterface(BaseInterface): + supports_monitor = True + def __init__(self, config): self.ctxt = Context('pva', nt=False) if 'EPICS_PVA_NAME_SERVERS' in os.environ: @@ -85,11 +87,20 @@ def wrapped_handler(value): return wrapped_handler - def monitor(self, handler, **kwargs): - for pv in self.in_list: + def monitor(self, handler, pv_name=None, **kwargs): + pvs = [pv_name] if pv_name else self.in_list + for pv in pvs: try: - new_handler = self.__handler_wrapper(handler, pv) - self.ctxt.monitor(pv, new_handler) + if pv_name: + # Per-PV handler (event-driven mode): extract raw value + def _wrap(h, p=pv): + def wrapped(value): + h(value['value']) + return wrapped + self.ctxt.monitor(pv, _wrap(handler)) + else: + new_handler = self.__handler_wrapper(handler, pv) + self.ctxt.monitor(pv, new_handler) except Exception as e: logger.error( f'Error monitoring in function monitor for SimplePVAInterface: {e}' @@ -98,20 +109,26 @@ def monitor(self, handler, **kwargs): raise e def get(self, name, **kwargs): - value = self.ctxt.get(name) - if isinstance(value['value'], np.ndarray): + raw = self.ctxt.get(name) + if isinstance(raw['value'], np.ndarray): # if value has dimension - if 'dimension' in value: - y_size = value['dimension'][0]['size'] - x_size = value['dimension'][1]['size'] - value = value['value'].reshape((y_size, x_size)) + if 'dimension' in raw: + y_size = raw['dimension'][0]['size'] + x_size = raw['dimension'][1]['size'] + value = raw['value'].reshape((y_size, x_size)) else: - value = value['value'] + value = raw['value'] else: - value = value['value'] + value = raw['value'] - value = {'value': value} - return name, value + result = {'value': value} + # Extract timestamp from p4p.Value for tracing + try: + ts = raw['timeStamp'] + result['timestamp'] = float(ts['secondsPastEpoch']) + float(ts['nanoseconds']) * 1e-9 + except Exception: + pass + return name, result @staticmethod def _payload_has_explicit_alarm(payload: Any) -> bool: @@ -240,19 +257,26 @@ def get_many(self, data, **kwargs): results = self.ctxt.get(data, throw=False) output = {} # print(f"results: {results}") - for value, key in zip(results, data): - if isinstance(value['value'], np.ndarray): + for raw, key in zip(results, data): + if isinstance(raw['value'], np.ndarray): # if value has dimension - if 'dimension' in value: - y_size = value['dimension'][0]['size'] - x_size = value['dimension'][1]['size'] - value = value['value'].reshape((y_size, x_size)) + if 'dimension' in raw: + y_size = raw['dimension'][0]['size'] + x_size = raw['dimension'][1]['size'] + value = raw['value'].reshape((y_size, x_size)) else: - value = value['value'] + value = raw['value'] else: - value = value['value'] + value = raw['value'] - output[key] = {'value': value} + result = {'value': value} + # Extract timestamp from p4p.Value for tracing + try: + ts = raw['timeStamp'] + result['timestamp'] = float(ts['secondsPastEpoch']) + float(ts['nanoseconds']) * 1e-9 + except Exception: + pass + output[key] = result return output @@ -275,6 +299,7 @@ class SimplePVAInterfaceServer(SimplePVAInterface): def __init__(self, config): super().__init__(config) self.shared_pvs = {} + self._pv_handlers = {} self.kv_store = {} if 'port' in config: @@ -347,6 +372,7 @@ def __init__(self, owner, pv_name: str, read_only: bool = False): self.owner = owner self.pv_name = pv_name self.read_only = read_only + self._monitor_callbacks = [] def put(self, pv: SharedPV, op: ServOpWrap): if self.read_only: @@ -357,21 +383,37 @@ def put(self, pv: SharedPV, op: ServOpWrap): self.pv_name, op.value() ) except Exception as exc: + logger.error(f'Handler.put prepare_write_payload error for {self.pv_name}: {exc}') + op.done(error=str(exc)) + return + try: + pv.post(payload, timestamp=time.time()) + except Exception as exc: + logger.error(f'Handler.put pv.post error for {self.pv_name}: {exc}') op.done(error=str(exc)) return - pv.post(payload, timestamp=time.time()) op.done() + logger.debug(f'Handler.put completed for {self.pv_name}, payload type={type(payload).__name__}, callbacks={len(self._monitor_callbacks)}') + # Notify event-driven monitor callbacks + raw_value = payload.get('value', payload) if isinstance(payload, dict) else payload + for cb in self._monitor_callbacks: + try: + cb(raw_value) + except Exception as e: + logger.error(f'Monitor callback error for {self.pv_name}: {e}') # PVs that are exclusively outputs are considered read-only read_only = False if 'mode' in pv_cfg: read_only = pv_cfg['mode'] == 'out' + h = Handler(self, pv, read_only) + self._pv_handlers[pv] = h pv_item = { pv: SharedPV( initial=pv_type_init, nt=pv_type_nt, - handler=Handler(self, pv, read_only), + handler=h, ) } # print(f"pv_item: {pv_item}") @@ -399,6 +441,16 @@ def close(self): self.server.stop() super().close() + def monitor(self, handler, pv_name=None, **kwargs): + """Register monitor callbacks on SharedPV handlers for event-driven mode.""" + pvs = [pv_name] if pv_name else self.in_list + for pv in pvs: + if pv not in self._pv_handlers: + logger.warning(f'PV {pv} not found in handlers, skipping monitor') + continue + self._pv_handlers[pv]._monitor_callbacks.append(handler) + logger.debug(f'Registered event-driven monitor callback for {pv}') + def put(self, name, value, **kwargs): payload, _ = self._prepare_write_payload(name, value) # if not open then open diff --git a/poly_lithic/src/utils/builder.py b/poly_lithic/src/utils/builder.py index 54f41cd..9bd1049 100644 --- a/poly_lithic/src/utils/builder.py +++ b/poly_lithic/src/utils/builder.py @@ -40,15 +40,24 @@ def __initailise_config(self, config_path): logger.error(f'Error initializing configuration: {e}') raise e - def build(self) -> MessageBroker: + def build(self, trace_store=None) -> MessageBroker: """Build the model manager.""" self.__build_observers() - self.__build_broker() + self.__build_broker(trace_store=trace_store) for name, observer in self.loaded_observers.items(): self.broker.attach(observer, self.config.modules[name].sub) + # Validate event-driven compatibility + if self.config.deployment.type == 'event_driven': + for obs in self.get_input_interface_observers(): + if not obs.interface.supports_monitor: + raise ValueError( + f'Interface {obs.interface.__class__.__name__} does not support ' + f'event-driven mode (supports_monitor=False)' + ) + return self.broker def __build_observers(self): @@ -128,7 +137,18 @@ def __build_observers(self): self.loaded_observers = loaded_observers return None - def __build_broker(self): + def __build_broker(self, trace_store=None): """Build the message broker.""" - self.broker = MessageBroker() + self.broker = MessageBroker(trace_store=trace_store) logger.debug(f'Built broker: {self.broker}') + + def get_input_interface_observers(self) -> list[InterfaceObserver]: + """Return InterfaceObservers subscribed to 'get_all' (input interfaces).""" + input_observers = [] + for name, observer in self.loaded_observers.items(): + if isinstance(observer, InterfaceObserver): + sub = self.config.modules[name].sub + subs = [sub] if isinstance(sub, str) else (sub or []) + if 'get_all' in subs: + input_observers.append(observer) + return input_observers diff --git a/poly_lithic/src/utils/messaging.py b/poly_lithic/src/utils/messaging.py index c6c27e2..12c4fe9 100644 --- a/poly_lithic/src/utils/messaging.py +++ b/poly_lithic/src/utils/messaging.py @@ -8,6 +8,9 @@ computed_field, ) import time +import threading +import collections +from uuid import uuid4 from poly_lithic.src.logging_utils import get_logger from poly_lithic.src.transformers import BaseTransformer from poly_lithic.src.interfaces import BaseInterface @@ -61,6 +64,9 @@ class Message(BaseModel): timestamp: float = Field(default_factory=time.time) # optional allow_unsafe: Optional[bool] = False + # tracing fields + trace_id: str = Field(default_factory=lambda: str(uuid4())) + parent_trace_ids: list[str] = Field(default_factory=list) @field_validator('topic') @classmethod @@ -146,13 +152,15 @@ def update(self, message: Message) -> Message: class MessageBroker: - def __init__(self): + def __init__(self, trace_store=None): """initialize the message broker""" self._observers: Dict[str, list[Observer]] = {} self._stats = {} self._stats_cnt = {} - self.queue = [] + self.queue = collections.deque() + self._queue_lock = threading.Lock() self.last_update = time.time() + self.trace_store = trace_store def attach(self, observer: Observer, topic: str | list[str]) -> None: """add observer to topic""" @@ -182,6 +190,9 @@ def detach(self, observer: Observer, topic: str | list[str]) -> None: # @profileit def notify(self, message: Message) -> None: """notify all observers of a message""" + if self.trace_store is not None: + self.trace_store.record(message) + if message.topic in self._observers: # logger.debug(f"notifying observers of {message}") @@ -200,10 +211,12 @@ def notify(self, message: Message) -> None: if result is not None: # if list of messages if isinstance(result, list): - for r in result: - self.queue.append(r) + with self._queue_lock: + for r in result: + self.queue.append(r) else: - self.queue.append(result) + with self._queue_lock: + self.queue.append(result) if time.time() - self.last_update > 1: self.last_update = time.time() @@ -237,12 +250,12 @@ def get_all(self) -> None: def parse_queue(self): """parse the queue and notify observers of each message""" - queue_snapshot = self.queue.copy() + with self._queue_lock: + queue_snapshot = list(self.queue) + self.queue.clear() for message in queue_snapshot: self.notify(message) - self.queue.remove(message) logger.debug(f'queue length: {len(self.queue)}') - # logger.debug(f"queue: {self.queue}") class TransformerObserver(Observer): @@ -255,6 +268,12 @@ def __init__( self.unpack_output = unpack_output def update(self, message: Message) -> Message | list[Message]: + # Snapshot input metadata before transformer strips it + input_meta = {} + for k, v in message.value.items(): + if isinstance(v, dict) and 'metadata' in v: + input_meta[k] = v['metadata'] + for key, value in message.value.items(): self.transformer.handler(key, value) @@ -266,9 +285,25 @@ def update(self, message: Message) -> Message | list[Message]: message_dict[key] = value else: message_dict[key] = {'value': value} + # Re-attach input metadata if this key had it + if key in input_meta: + existing_meta = message_dict[key].get('metadata', {}) + existing_meta.update(input_meta[key]) + message_dict[key]['metadata'] = existing_meta self.transformer.updated = False - return Message(topic=self.topic, source=str(self), value=message_dict) + out_msg = Message( + topic=self.topic, + source=str(self), + value=message_dict, + parent_trace_ids=[message.trace_id], + ) + # Inject trace_id into each variable struct metadata + for key in out_msg.value: + if isinstance(out_msg.value[key], dict): + meta = out_msg.value[key].setdefault('metadata', {}) + meta['trace'] = {'trace_id': out_msg.trace_id} + return out_msg class InterfaceObserver(Observer): @@ -309,7 +344,15 @@ def update(self, message: Message) -> Message | list[Message]: logger.debug(f'updating {self}') if os.environ['PUBLISH'] == 'True': - self.interface.put_many(message.value) + # Strip internal metadata before writing to the interface + # (e.g. p4p NTScalar doesn't have a 'metadata' field) + clean = {} + for k, v in message.value.items(): + if isinstance(v, dict) and 'metadata' in v: + clean[k] = {fk: fv for fk, fv in v.items() if fk != 'metadata'} + else: + clean[k] = v + self.interface.put_many(clean) else: logger.warning( 'PUBLISH is set to False, this will not publish to the interface' @@ -337,7 +380,13 @@ def get_all(self) -> list[Message]: if value is not None: output_dict[key] = value - messages.append(Message(topic=self.topic, source=str(self), value=output_dict)) + msg = Message(topic=self.topic, source=str(self), value=output_dict) + # Inject trace_id into each variable struct metadata (origin — no parent) + for key in msg.value: + if isinstance(msg.value[key], dict): + meta = msg.value[key].setdefault('metadata', {}) + meta['trace'] = {'trace_id': msg.trace_id} + messages.append(msg) return messages # if self.last_get_all is not None: @@ -381,6 +430,47 @@ def put_many(self, message: Message) -> None: raise ValueError('message value must be a dictionary') self.interface.put_many(message.value) + def start_monitors(self, broker, min_interval: float = 0.0, on_change_only: bool = False): + """Start PV monitors that push updates into the broker queue (event-driven mode). + + Args: + broker: MessageBroker instance to push messages into. + min_interval: Minimum seconds between updates per PV (throttle). + on_change_only: When True, skip updates where the value hasn't changed. + """ + last_fire_time = {} + last_value = {} + self._monitors = [] + + def _make_handler(pv_name): + def handler(value): + now = time.time() + # Throttle + if min_interval > 0: + if pv_name in last_fire_time and (now - last_fire_time[pv_name]) < min_interval: + return + # Dedup + if on_change_only: + if pv_name in last_value and last_value[pv_name] == value: + return + last_value[pv_name] = value + + last_fire_time[pv_name] = now + val_struct = {'value': value} + msg = Message( + topic=self.topic, + source=str(self), + value={pv_name: val_struct}, + ) + val_struct.setdefault('metadata', {})['trace'] = {'trace_id': msg.trace_id} + with broker._queue_lock: + broker.queue.append(msg) + return handler + + for pv_name in self.interface.get_inputs(): + mon = self.interface.monitor(_make_handler(pv_name), pv_name) + self._monitors.append(mon) + class MockModel: def __init__(self): @@ -470,6 +560,12 @@ def update(self, message: Message) -> list[Message]: messages = [] logger.debug(f'updating {self}') + # Snapshot input metadata before unpacking + input_meta = {} + for k, v in message.value.items(): + if isinstance(v, dict) and 'metadata' in v: + input_meta[k] = v['metadata'] + if self.unpack_input: # logger.debug(f"unpacking input: {message.value}") value = {v: message.value[v]['value'] for v in message.value} @@ -491,7 +587,18 @@ def update(self, message: Message) -> list[Message]: # logger.debug(f"not packing output passign raw: {pred}") output = pred - messages.append(Message(topic=self.topic, source=str(self), value=output)) + out_msg = Message( + topic=self.topic, + source=str(self), + value=output, + parent_trace_ids=[message.trace_id], + ) + # Inject trace_id and merge input metadata into each output variable struct + for key in out_msg.value: + if isinstance(out_msg.value[key], dict): + meta = out_msg.value[key].setdefault('metadata', {}) + meta['trace'] = {'trace_id': out_msg.trace_id} + messages.append(out_msg) return messages diff --git a/poly_lithic/src/utils/trace_server.py b/poly_lithic/src/utils/trace_server.py new file mode 100644 index 0000000..551942b --- /dev/null +++ b/poly_lithic/src/utils/trace_server.py @@ -0,0 +1,69 @@ +"""Standalone FastAPI server for querying message trace records.""" + +import threading +from dataclasses import asdict + +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import HTMLResponse + +from .trace_store import TraceStore + +_VIEWER_HTML = None + + +def _load_viewer_html(): + global _VIEWER_HTML + if _VIEWER_HTML is None: + import pathlib + p = pathlib.Path(__file__).with_name("trace_viewer.html") + _VIEWER_HTML = p.read_text() + return _VIEWER_HTML + + +def create_trace_app(trace_store: TraceStore) -> FastAPI: + """Create a FastAPI app wired to the given TraceStore.""" + app = FastAPI(title="Poly-Lithic Tracing API") + app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["GET"], allow_headers=["*"]) + + @app.get("/", response_class=HTMLResponse) + def viewer(): + return _load_viewer_html() + + @app.get("/health") + def health(): + return {"status": "ok"} + + @app.get("/traces") + def get_traces(limit: int = 100): + records = trace_store.get_recent(limit) + return [asdict(r) for r in records] + + @app.get("/traces/{trace_id}") + def get_trace(trace_id: str): + record = trace_store.get(trace_id) + if record is None: + raise HTTPException(status_code=404, detail="Trace not found") + return asdict(record) + + @app.get("/traces/{trace_id}/lineage") + def get_lineage(trace_id: str): + records = trace_store.get_lineage(trace_id) + if not records: + raise HTTPException(status_code=404, detail="Trace not found") + return [asdict(r) for r in records] + + return app + + +def start_trace_server(trace_store: TraceStore, port: int = 8100) -> threading.Thread: + """Start the tracing FastAPI server in a background daemon thread.""" + app = create_trace_app(trace_store) + + def _run(): + uvicorn.run(app, host="0.0.0.0", port=port, log_level="warning") + + thread = threading.Thread(target=_run, name="trace-server", daemon=True) + thread.start() + return thread diff --git a/poly_lithic/src/utils/trace_store.py b/poly_lithic/src/utils/trace_store.py new file mode 100644 index 0000000..efaa0b5 --- /dev/null +++ b/poly_lithic/src/utils/trace_store.py @@ -0,0 +1,107 @@ +"""In-memory ring buffer for message trace records.""" + +import collections +import threading +import time +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + + +def _serialise_value(v: Any) -> Any: + """Convert a value to something JSON-safe.""" + if isinstance(v, np.ndarray): + return v.tolist() + if isinstance(v, (np.integer,)): + return int(v) + if isinstance(v, (np.floating,)): + return float(v) + if isinstance(v, dict): + return {k: _serialise_value(val) for k, val in v.items()} + if isinstance(v, (list, tuple)): + return [_serialise_value(i) for i in v] + # Handle p4p scalar wrappers (ntfloat, ntint, etc.) which are float/int + # subclasses but contain unpicklable C-level state. + if isinstance(v, float): + return float(v) + if isinstance(v, int) and not isinstance(v, bool): + return int(v) + return v + + +@dataclass +class TraceRecord: + trace_id: str + parent_trace_ids: list[str] + topic: str + source: str + timestamp: float + variable_keys: list[str] = field(default_factory=list) + variable_values: dict[str, Any] = field(default_factory=dict) + + +class TraceStore: + """Thread-safe in-memory ring buffer of TraceRecords.""" + + def __init__(self, maxlen: int = 10000): + self._buffer: collections.deque[TraceRecord] = collections.deque(maxlen=maxlen) + self._index: dict[str, TraceRecord] = {} + self._lock = threading.Lock() + self._maxlen = maxlen + + def record(self, message) -> None: + """Record a message as a TraceRecord.""" + vals = {} + if isinstance(message.value, dict): + for k, v in message.value.items(): + if isinstance(v, dict) and 'value' in v: + vals[k] = _serialise_value(v['value']) + else: + vals[k] = _serialise_value(v) + rec = TraceRecord( + trace_id=message.trace_id, + parent_trace_ids=list(message.parent_trace_ids), + topic=message.topic, + source=message.source, + timestamp=time.time(), + variable_keys=list(message.value.keys()) if isinstance(message.value, dict) else [], + variable_values=vals, + ) + with self._lock: + # If buffer is full, the evicted record should be removed from index + if len(self._buffer) == self._maxlen: + evicted = self._buffer[0] + self._index.pop(evicted.trace_id, None) + self._buffer.append(rec) + self._index[rec.trace_id] = rec + + def get(self, trace_id: str) -> TraceRecord | None: + """Get a single trace record by ID.""" + with self._lock: + return self._index.get(trace_id) + + def get_lineage(self, trace_id: str) -> list[TraceRecord]: + """Walk parent_trace_ids to build the full ancestry chain.""" + with self._lock: + result = [] + visited = set() + queue = [trace_id] + while queue: + tid = queue.pop(0) + if tid in visited: + continue + visited.add(tid) + rec = self._index.get(tid) + if rec is not None: + result.append(rec) + for pid in rec.parent_trace_ids: + if pid not in visited: + queue.append(pid) + return result + + def get_recent(self, limit: int = 100) -> list[TraceRecord]: + """Return the most recent trace records.""" + with self._lock: + items = list(self._buffer) + return items[-limit:] diff --git a/poly_lithic/src/utils/trace_viewer.html b/poly_lithic/src/utils/trace_viewer.html new file mode 100644 index 0000000..5e98621 --- /dev/null +++ b/poly_lithic/src/utils/trace_viewer.html @@ -0,0 +1,450 @@ + + + + + +Poly-Lithic Trace Viewer + + + + +
+

▸ Trace Viewer

+
+ + : + + + + + + +
+
+ +
+
+
+ +
+ + + +
TimeTrace IDTopicSourceValuesParents
+
Enter server address and click Fetch
+
+
+
Select a trace to view details
+ +
+
+ + + + diff --git a/tests/test_event_driven.py b/tests/test_event_driven.py new file mode 100644 index 0000000..b5139a7 --- /dev/null +++ b/tests/test_event_driven.py @@ -0,0 +1,917 @@ +# Tests for the event-driven deployment mode. +# Each test isolates one step in the chain to find exactly where it breaks. + +import socket +import time +import logging +import pytest +from unittest.mock import MagicMock + +from p4p.client.thread import Context + +from poly_lithic.src.utils.messaging import ( + Message, + MessageBroker, + InterfaceObserver, + TransformerObserver, + ModelObserver, +) +from poly_lithic.src.transformers.BaseTransformers import SimpleTransformer +from poly_lithic.src.interfaces import registered_interfaces + + +def _get_free_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(('127.0.0.1', 0)) + return sock.getsockname()[1] + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def p4p_server(): + config = { + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 2.0}, + 'B1': {'name': 'B1', 'proto': 'pva', 'type': 'scalar', 'default': 1.0}, + } + } + srv = registered_interfaces['p4p_server'](config) + yield srv + srv.close() + + +@pytest.fixture +def interface_observer(p4p_server): + return InterfaceObserver(p4p_server, 'in_interface') + + +@pytest.fixture +def transformer_in(): + config = { + 'variables': {'x2': {'formula': 'A1 + B1'}, 'x1': {'formula': 'A1'}}, + 'symbols': ['A1', 'B1'], + } + return SimpleTransformer(config) + + +@pytest.fixture +def transformer_observer_in(transformer_in): + return TransformerObserver(transformer_in, 'in_transformer') + + +class MockModel: + def __init__(self): + self.call_log = [] + + def evaluate(self, value): + self.call_log.append(value) + return {'pred0': value['x1'] + value['x2']} + + +@pytest.fixture +def mock_model(): + return MockModel() + + +@pytest.fixture +def model_observer(mock_model): + return ModelObserver(model=mock_model, topic='model_out') + + +@pytest.fixture +def broker(): + return MessageBroker() + + +# --------------------------------------------------------------------------- +# Step 1: Handler.put() fires monitor callbacks +# --------------------------------------------------------------------------- + +class TestHandlerMonitorCallbacks: + """Verify that when an external client writes to a SharedPV, + the Handler.put() method fires registered monitor callbacks.""" + + def test_handler_has_callback_list(self, p4p_server): + """Each PV handler should have a _monitor_callbacks list.""" + for pv_name, handler in p4p_server._pv_handlers.items(): + assert hasattr(handler, '_monitor_callbacks'), ( + f'Handler for {pv_name} missing _monitor_callbacks' + ) + + def test_monitor_registers_callback(self, p4p_server): + """Calling monitor(handler, pv_name) should register a callback.""" + cb = MagicMock() + p4p_server.monitor(cb, pv_name='A1') + assert cb in p4p_server._pv_handlers['A1']._monitor_callbacks + + def test_handler_put_fires_callbacks(self, p4p_server): + """Simulate what Handler.put() does: post + fire callbacks.""" + cb = MagicMock() + p4p_server.monitor(cb, pv_name='A1') + + # Simulate an external put by writing via the server's own put method + # (this bypasses the Handler — we need to test the Handler directly) + handler = p4p_server._pv_handlers['A1'] + + # Directly invoke the callback path with a known value + raw_value = 42.0 + for c in handler._monitor_callbacks: + c(raw_value) + + cb.assert_called_once_with(42.0) + + def test_server_put_does_not_fire_monitor(self, p4p_server): + """Server's own put() should NOT fire monitor callbacks. + Only external client writes (Handler.put) should trigger them.""" + cb = MagicMock() + p4p_server.monitor(cb, pv_name='A1') + + # Use the server's put — this goes through SharedPV.post, not Handler.put + p4p_server.put('A1', {'value': 99.0}) + + # This should NOT have fired because put() bypasses Handler.put() + cb.assert_not_called() + + +# --------------------------------------------------------------------------- +# Step 2: start_monitors creates handlers that produce Messages +# --------------------------------------------------------------------------- + +class TestStartMonitors: + """Verify that start_monitors registers per-PV handlers that + create Messages and push them into the broker queue.""" + + def test_start_monitors_registers_callbacks(self, interface_observer, broker): + interface_observer.start_monitors(broker) + # Should have registered monitors for each input PV + assert len(interface_observer._monitors) > 0 + + def test_monitor_handler_creates_message(self, interface_observer, broker): + """Manually invoke the handler created by start_monitors and check + that a Message appears in the broker queue.""" + interface_observer.start_monitors(broker) + + # Directly fire a callback on one PV handler + pv_name = 'A1' + handler = interface_observer.interface._pv_handlers[pv_name] + assert len(handler._monitor_callbacks) > 0, ( + 'start_monitors did not register any callbacks on Handler' + ) + + # Simulate the callback being fired with a scalar value + for cb in handler._monitor_callbacks: + cb(5.0) + + assert len(broker.queue) == 1, f'Expected 1 message in queue, got {len(broker.queue)}' + msg = broker.queue[0] + assert msg.topic == 'in_interface' + assert 'A1' in msg.value + assert msg.value['A1']['value'] == 5.0 + + def test_monitor_handler_message_value_structure(self, interface_observer, broker): + """The message value should be {pv_name: {'value': scalar}} + which is what TransformerObserver.update() expects.""" + interface_observer.start_monitors(broker) + + handler = interface_observer.interface._pv_handlers['A1'] + for cb in handler._monitor_callbacks: + cb(7.5) + + msg = broker.queue[0] + val = msg.value['A1'] + assert isinstance(val, dict), f'Expected dict, got {type(val)}' + assert 'value' in val, f'Missing "value" key in {val}' + assert isinstance(val['value'], (int, float)), ( + f'Expected numeric value, got {type(val["value"])}: {val["value"]}' + ) + + +# --------------------------------------------------------------------------- +# Step 3: TransformerObserver handles single-PV messages +# --------------------------------------------------------------------------- + +class TestTransformerSinglePV: + """Verify that the transformer handles single-PV event messages + and produces output once all inputs have been seen at least once.""" + + def test_single_pv_does_not_transform_until_all_seen(self, transformer_in): + """Sending only one PV should NOT trigger transform.""" + transformer_in.handler('A1', {'value': 5.0}) + assert not transformer_in.updated, ( + 'Transformer should not update with only one input set' + ) + + def test_both_pvs_triggers_transform(self, transformer_in): + """After both PVs have values, transform should fire.""" + transformer_in.handler('A1', {'value': 5.0}) + transformer_in.handler('B1', {'value': 3.0}) + assert transformer_in.updated, ( + 'Transformer should update once all inputs are set' + ) + assert 'x1' in transformer_in.latest_transformed + assert 'x2' in transformer_in.latest_transformed + + def test_subsequent_single_pv_triggers_transform(self, transformer_in): + """Once all inputs have been initialized, a single PV update + should trigger transform because all() is already satisfied.""" + # Initialize both + transformer_in.handler('A1', {'value': 5.0}) + transformer_in.handler('B1', {'value': 3.0}) + assert transformer_in.updated + transformer_in.updated = False # reset + + # Now only one changes + transformer_in.handler('A1', {'value': 10.0}) + assert transformer_in.updated, ( + 'Transformer should re-fire on single-PV update after initialization' + ) + + def test_transformer_observer_with_single_pv_message( + self, transformer_observer_in, transformer_in + ): + """TransformerObserver.update() with a single-PV message.""" + # First seed both PVs so all() passes + msg_a = Message(topic='in_interface', source='test', value={ + 'A1': {'value': 5.0}, + 'B1': {'value': 3.0}, + }) + result = transformer_observer_in.update(msg_a) + assert result is not None, 'First full message should produce output' + + # Now send a single-PV update + msg_b = Message(topic='in_interface', source='test', value={ + 'A1': {'value': 10.0}, + }) + result = transformer_observer_in.update(msg_b) + assert result is not None, ( + 'Single-PV update after seeding should produce output' + ) + + +# --------------------------------------------------------------------------- +# Step 4: Full chain — monitor → transformer → model +# --------------------------------------------------------------------------- + +class TestFullEventDrivenChain: + """End-to-end test: simulate monitor callbacks flowing through + the full broker pipeline.""" + + def test_full_chain_after_seeding( + self, + interface_observer, + transformer_observer_in, + model_observer, + mock_model, + broker, + ): + """After seeding with get_all, a single PV monitor event should + flow through the entire chain.""" + # Wire up broker + broker.attach(interface_observer, 'get_all') + broker.attach(transformer_observer_in, 'in_interface') + broker.attach(model_observer, 'in_transformer') + + # Seed: get_all populates both PVs + broker.get_all() + # parse until queue drains (get_all → interface → transformer → model) + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + assert len(mock_model.call_log) == 1, ( + f'Model should have been called once from get_all, ' + f'got {len(mock_model.call_log)}' + ) + + # Now start monitors + interface_observer.start_monitors(broker) + + # Simulate external PV write + handler = interface_observer.interface._pv_handlers['A1'] + for cb in handler._monitor_callbacks: + cb(99.0) + + assert len(broker.queue) == 1 + # Parse through the chain + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + assert len(mock_model.call_log) == 2, ( + f'Model should have been called a second time from monitor event, ' + f'got {len(mock_model.call_log)}' + ) + + def test_chain_without_seeding_requires_all_pvs( + self, + interface_observer, + transformer_observer_in, + model_observer, + mock_model, + broker, + ): + """Without get_all seeding, the model should only fire once + ALL input PVs have been written via monitors.""" + broker.attach(transformer_observer_in, 'in_interface') + broker.attach(model_observer, 'in_transformer') + + # Start monitors (no get_all) + interface_observer.start_monitors(broker) + + # Write only A1 + handler_a = interface_observer.interface._pv_handlers['A1'] + for cb in handler_a._monitor_callbacks: + cb(5.0) + + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + assert len(mock_model.call_log) == 0, ( + 'Model should NOT fire with only one PV written' + ) + + # Now write B1 + handler_b = interface_observer.interface._pv_handlers['B1'] + for cb in handler_b._monitor_callbacks: + cb(3.0) + + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + assert len(mock_model.call_log) == 1, ( + f'Model should fire once both PVs are written, ' + f'got {len(mock_model.call_log)}' + ) + + +# --------------------------------------------------------------------------- +# Step 5: Real p4p client write triggers Handler.put → monitor callback +# --------------------------------------------------------------------------- + +class TestRealClientWrite: + """Use a real p4p Context client to write to the server and verify + that Handler.put() fires monitor callbacks end-to-end.""" + + def test_client_put_fires_monitor_callback(self, monkeypatch): + """A p4p client writing to a server PV should trigger the + Handler.put() path which fires monitor callbacks.""" + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + }, + } + server = registered_interfaces['p4p_server'](server_config) + try: + cb = MagicMock() + server.monitor(cb, pv_name='A1') + + # Write via a real p4p client + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('A1', {'value': 42.0}) + time.sleep(0.3) # give Handler.put a moment to fire + finally: + client.close() + + assert cb.call_count >= 1, ( + f'Monitor callback should have been called by client put, ' + f'but was called {cb.call_count} times' + ) + # Check the value passed to the callback + call_args = cb.call_args[0] + assert call_args[0] == pytest.approx(42.0), ( + f'Callback received {call_args[0]}, expected 42.0' + ) + # CRITICAL: Check the TYPE — must be a scalar, not p4p.Value + assert isinstance(call_args[0], (int, float)), ( + f'Callback value should be a scalar (int/float), ' + f'got {type(call_args[0])}: {call_args[0]}' + ) + finally: + server.close() + + def test_client_put_full_chain(self, monkeypatch): + """Full chain: real client put → Handler.put → monitor callback → + broker queue → transformer → model.""" + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + 'B1': {'name': 'B1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + }, + } + server = registered_interfaces['p4p_server'](server_config) + broker = MessageBroker() + obs = InterfaceObserver(server, 'in_interface') + + config_t = { + 'variables': {'x2': {'formula': 'A1 + B1'}, 'x1': {'formula': 'A1'}}, + 'symbols': ['A1', 'B1'], + } + t_obs = TransformerObserver(SimpleTransformer(config_t), 'in_transformer') + + mock_model = MockModel() + m_obs = ModelObserver(model=mock_model, topic='model_out') + + broker.attach(obs, 'get_all') + broker.attach(t_obs, 'in_interface') + broker.attach(m_obs, 'in_transformer') + + # Seed via get_all so transformer has all inputs + broker.get_all() + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + initial_calls = len(mock_model.call_log) + + # Start monitors + obs.start_monitors(broker) + + # Write via a real p4p client + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('A1', {'value': 99.0}) + time.sleep(0.3) + finally: + client.close() + + # Parse the queue + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + assert len(mock_model.call_log) > initial_calls, ( + f'Model should have been called after client put, ' + f'call_log has {len(mock_model.call_log)} entries ' + f'(initial: {initial_calls})' + ) + + server.close() + + def test_client_put_with_colon_pv_names(self, monkeypatch): + """Test with PV names containing colons (like ML:LOCAL:TEST_A) + since that's what the real deployment config uses.""" + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'ML:LOCAL:TEST_A': { + 'name': 'ML:LOCAL:TEST_A', + 'proto': 'pva', + 'type': 'scalar', + 'default': 0.0, + }, + 'ML:LOCAL:TEST_B': { + 'name': 'ML:LOCAL:TEST_B', + 'proto': 'pva', + 'type': 'scalar', + 'default': 0.0, + }, + 'ML:LOCAL:TEST_S': { + 'name': 'ML:LOCAL:TEST_S', + 'proto': 'pva', + 'type': 'scalar', + 'default': 0.0, + }, + }, + } + server = registered_interfaces['p4p_server'](server_config) + broker = MessageBroker() + obs = InterfaceObserver(server, 'in_interface') + + config_t = { + 'variables': { + 'x': {'formula': 'ML:LOCAL:TEST_A * 2 + 10'}, + 'y': {'formula': 'ML:LOCAL:TEST_B + 120'}, + }, + 'symbols': ['ML:LOCAL:TEST_B', 'ML:LOCAL:TEST_A'], + } + t_obs = TransformerObserver(SimpleTransformer(config_t), 'in_transformer') + + class CoordModel: + def __init__(self): + self.call_log = [] + + def evaluate(self, value): + self.call_log.append(value) + return {'output': value['x'] + value['y']} + + mock_model = CoordModel() + m_obs = ModelObserver(model=mock_model, topic='model_out') + + config_t_out = { + 'variables': {'ML:LOCAL:TEST_S': {'formula': 'output'}}, + 'symbols': ['output'], + } + t_obs_out = TransformerObserver( + SimpleTransformer(config_t_out), 'out_transformer', unpack_output=True + ) + + import os + os.environ['PUBLISH'] = 'True' + obs_out = InterfaceObserver(server, 'out_interface') + + broker.attach(obs, 'get_all') + broker.attach(t_obs, 'in_interface') + broker.attach(m_obs, 'in_transformer') + broker.attach(t_obs_out, 'model_out') + broker.attach(obs_out, 'out_transformer') + + # Seed + broker.get_all() + for _ in range(20): + if not broker.queue: + break + broker.parse_queue() + + initial_calls = len(mock_model.call_log) + assert initial_calls == 1, f'Expected 1 initial model call, got {initial_calls}' + + # Start monitors + obs.start_monitors(broker) + + # Write via real client + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('ML:LOCAL:TEST_A', {'value': 5.0}) + time.sleep(0.3) + finally: + client.close() + + for _ in range(20): + if not broker.queue: + break + broker.parse_queue() + + assert len(mock_model.call_log) > initial_calls, ( + f'Model should have been called after client put to ML:LOCAL:TEST_A, ' + f'call_log={mock_model.call_log}' + ) + + # Check the output PV was updated + _, result = server.get('ML:LOCAL:TEST_S') + logging.info(f'ML:LOCAL:TEST_S value after event: {result}') + + server.close() + + def test_client_put_both_pvs_no_seeding(self, monkeypatch): + """Without get_all seeding, writing both PVs via real client + should trigger the model (the user's exact scenario).""" + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'ML:LOCAL:TEST_A': { + 'name': 'ML:LOCAL:TEST_A', + 'proto': 'pva', + 'type': 'scalar', + 'default': 0.0, + }, + 'ML:LOCAL:TEST_B': { + 'name': 'ML:LOCAL:TEST_B', + 'proto': 'pva', + 'type': 'scalar', + 'default': 0.0, + }, + 'ML:LOCAL:TEST_S': { + 'name': 'ML:LOCAL:TEST_S', + 'proto': 'pva', + 'type': 'scalar', + 'default': 0.0, + }, + }, + } + server = registered_interfaces['p4p_server'](server_config) + broker = MessageBroker() + obs = InterfaceObserver(server, 'in_interface') + + config_t = { + 'variables': { + 'x': {'formula': 'ML:LOCAL:TEST_A * 2 + 10'}, + 'y': {'formula': 'ML:LOCAL:TEST_B + 120'}, + }, + 'symbols': ['ML:LOCAL:TEST_B', 'ML:LOCAL:TEST_A'], + } + t_obs = TransformerObserver(SimpleTransformer(config_t), 'in_transformer') + + class CoordModel2: + def __init__(self): + self.call_log = [] + + def evaluate(self, value): + self.call_log.append(value) + return {'output': value['x'] + value['y']} + + mock_model = CoordModel2() + m_obs = ModelObserver(model=mock_model, topic='model_out') + + # NO get_all subscription — just transformer and model + broker.attach(t_obs, 'in_interface') + broker.attach(m_obs, 'in_transformer') + + # Start monitors (no seeding!) + obs.start_monitors(broker) + + # Write BOTH PVs via real client (user's exact scenario) + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('ML:LOCAL:TEST_A', {'value': 5.0}) + time.sleep(0.2) + + # Parse after first write + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + logging.info(f'After first put: model calls={len(mock_model.call_log)}') + + client.put('ML:LOCAL:TEST_B', {'value': 3.0}) + time.sleep(0.2) + + # Parse after second write + for _ in range(10): + if not broker.queue: + break + broker.parse_queue() + + logging.info(f'After second put: model calls={len(mock_model.call_log)}') + finally: + client.close() + + assert len(mock_model.call_log) >= 1, ( + f'Model should have been called after both PVs written, ' + f'call_log={mock_model.call_log}' + ) + + server.close() + + +# --------------------------------------------------------------------------- +# Step 6: Mimic cli.py event_driven loop with asyncio +# --------------------------------------------------------------------------- + +class TestCliEventLoop: + """Replicate what cli.py model_main does in event_driven mode to + find if the issue is in the async loop or startup sequence.""" + + @pytest.mark.asyncio + async def test_asyncio_event_loop_with_seeding(self, monkeypatch): + """Mimic cli.py event_driven mode WITH get_all seeding.""" + import asyncio + + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + 'B1': {'name': 'B1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + }, + } + server = registered_interfaces['p4p_server'](server_config) + broker = MessageBroker() + obs = InterfaceObserver(server, 'in_interface') + + config_t = { + 'variables': {'x2': {'formula': 'A1 + B1'}, 'x1': {'formula': 'A1'}}, + 'symbols': ['A1', 'B1'], + } + t_obs = TransformerObserver(SimpleTransformer(config_t), 'in_transformer') + + mock_model = MockModel() + m_obs = ModelObserver(model=mock_model, topic='model_out') + + broker.attach(obs, 'get_all') + broker.attach(t_obs, 'in_interface') + broker.attach(m_obs, 'in_transformer') + + # Seed with get_all + broker.get_all() + while broker.queue: + broker.parse_queue() + + initial_calls = len(mock_model.call_log) + + # Start monitors + obs.start_monitors(broker) + + # Write via real client + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('A1', {'value': 55.0}) + finally: + client.close() + + # Run async loop like cli.py + for _ in range(100): + if broker.queue: + broker.parse_queue() + if len(mock_model.call_log) > initial_calls: + break + await asyncio.sleep(0.01) + + assert len(mock_model.call_log) > initial_calls, ( + f'Model should fire in async loop after seeding, ' + f'got {len(mock_model.call_log)} calls (initial: {initial_calls})' + ) + server.close() + + @pytest.mark.asyncio + async def test_asyncio_event_loop_without_seeding(self, monkeypatch): + """Mimic cli.py WITHOUT seeding — the current broken behavior.""" + import asyncio + + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + 'B1': {'name': 'B1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + }, + } + server = registered_interfaces['p4p_server'](server_config) + broker = MessageBroker() + obs = InterfaceObserver(server, 'in_interface') + + config_t = { + 'variables': {'x2': {'formula': 'A1 + B1'}, 'x1': {'formula': 'A1'}}, + 'symbols': ['A1', 'B1'], + } + t_obs = TransformerObserver(SimpleTransformer(config_t), 'in_transformer') + + mock_model = MockModel() + m_obs = ModelObserver(model=mock_model, topic='model_out') + + broker.attach(t_obs, 'in_interface') + broker.attach(m_obs, 'in_transformer') + + # Start monitors (no seeding!) + obs.start_monitors(broker) + + # Write both PVs via real client + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('A1', {'value': 5.0}) + for _ in range(30): + if broker.queue: + broker.parse_queue() + await asyncio.sleep(0.01) + + client.put('B1', {'value': 3.0}) + for _ in range(30): + if broker.queue: + broker.parse_queue() + await asyncio.sleep(0.01) + finally: + client.close() + + logging.info(f'No-seed async: model calls={len(mock_model.call_log)}') + assert len(mock_model.call_log) >= 1, ( + f'Model should fire after writing both PVs (no seeding), ' + f'call_log={mock_model.call_log}' + ) + server.close() + + +# --------------------------------------------------------------------------- +# Step 7: Diagnose exact value types flowing through the chain +# --------------------------------------------------------------------------- + +class TestValueTypeDiagnostics: + """Instrument every step to see the exact types and values.""" + + def test_handler_put_payload_type(self, monkeypatch): + """Instrument Handler.put() to log the exact payload type + and raw_value type passed to callbacks.""" + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + }, + } + server = registered_interfaces['p4p_server'](server_config) + + received_values = [] + received_types = [] + + def diagnostic_cb(value): + received_types.append(type(value).__name__) + received_values.append(repr(value)) + + server.monitor(diagnostic_cb, pv_name='A1') + + try: + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('A1', {'value': 42.0}) + time.sleep(0.3) + finally: + client.close() + + logging.warning(f'Received types: {received_types}') + logging.warning(f'Received values: {received_values}') + + assert len(received_types) >= 1, 'No callback received' + # ntfloat is p4p's scalar wrapper — it's a subclass of float + assert received_types[0] in ('float', 'int', 'float64', 'ntfloat'), ( + f'Expected scalar type, got {received_types[0]}: {received_values[0]}' + ) + finally: + server.close() + + def test_full_message_value_type_through_chain(self, monkeypatch): + """Trace the value type at each step: Handler.put -> queue -> transformer.""" + port = _get_free_port() + monkeypatch.setenv('EPICS_PVA_NAME_SERVERS', f'127.0.0.1:{port}') + + server_config = { + 'port': port, + 'variables': { + 'A1': {'name': 'A1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + 'B1': {'name': 'B1', 'proto': 'pva', 'type': 'scalar', 'default': 0.0}, + }, + } + server = registered_interfaces['p4p_server'](server_config) + broker = MessageBroker() + obs = InterfaceObserver(server, 'in_interface') + + config_t = { + 'variables': {'x2': {'formula': 'A1 + B1'}, 'x1': {'formula': 'A1'}}, + 'symbols': ['A1', 'B1'], + } + t_obs = TransformerObserver(SimpleTransformer(config_t), 'in_transformer') + mock_model = MockModel() + m_obs = ModelObserver(model=mock_model, topic='model_out') + + broker.attach(obs, 'get_all') + broker.attach(t_obs, 'in_interface') + broker.attach(m_obs, 'in_transformer') + + # Seed + broker.get_all() + while broker.queue: + broker.parse_queue() + + obs.start_monitors(broker) + + client = Context('pva', conf={'EPICS_PVA_NAME_SERVERS': f'127.0.0.1:{port}'}) + try: + client.put('A1', {'value': 99.0}) + time.sleep(0.3) + finally: + client.close() + + # Check what's in the queue before parsing + assert len(broker.queue) >= 1, 'No message in queue after client put' + msg = broker.queue[0] + logging.warning(f'Queue message topic: {msg.topic}') + logging.warning(f'Queue message value: {msg.value}') + for k, v in msg.value.items(): + logging.warning(f' key={k}, value_type={type(v)}, value={v}') + if isinstance(v, dict) and 'value' in v: + inner = v['value'] + logging.warning(f' inner_type={type(inner).__name__}, inner={inner}') + assert isinstance(inner, (int, float)), ( + f'Message value for {k} should be scalar, ' + f'got {type(inner).__name__}: {inner}' + ) + + # Parse and check model was called + while broker.queue: + broker.parse_queue() + + assert len(mock_model.call_log) >= 2, ( + f'Model should have been called, ' + f'call_log={mock_model.call_log}' + ) + server.close() diff --git a/tests/test_interfaceobserver.py b/tests/test_interfaceobserver.py index a134938..a8ea3bb 100644 --- a/tests/test_interfaceobserver.py +++ b/tests/test_interfaceobserver.py @@ -119,7 +119,9 @@ def test_interface_observer_get_all(interface_observer): assert len(result) == 1 assert result[0].topic == 'next_step' - assert result[0].value['test_scalar'] == {'value': 9.0} + assert result[0].value['test_scalar']['value'] == 9.0 + # trace metadata is injected by InterfaceObserver.get_all + assert 'trace' in result[0].value['test_scalar'].get('metadata', {}) np.testing.assert_array_equal( result[0].value['test_array']['value'], np.array([7.0, 8.0, 9.0]) From 8da73f70cc3d6e318c0712f04d548512d3d73485 Mon Sep 17 00:00:00 2001 From: MatInGit Date: Tue, 7 Apr 2026 13:58:27 +0100 Subject: [PATCH 2/4] Add tests for trace lineage completeness in TransformerObserver and ModelObserver - Implemented unit tests to verify that parent_trace_ids correctly aggregate trace_ids from all contributing inputs, not just the triggering message. - Created tests for SimpleTransformer and PassThroughTransformer to ensure lineage behavior is consistent. - Added tests for ModelObserver to check aggregation of trace_ids from input metadata. - Included a test for transformers without latest_input_struct to confirm they still function correctly. --- poly_lithic/src/utils/messaging.py | 31 +- poly_lithic/src/utils/readme.md | 1101 ++++++++++++++++++++++++++++ readme.md | 139 ++++ tests/test_trace_lineage.py | 280 +++++++ 4 files changed, 1548 insertions(+), 3 deletions(-) create mode 100644 poly_lithic/src/utils/readme.md create mode 100644 tests/test_trace_lineage.py diff --git a/poly_lithic/src/utils/messaging.py b/poly_lithic/src/utils/messaging.py index 12c4fe9..8badc32 100644 --- a/poly_lithic/src/utils/messaging.py +++ b/poly_lithic/src/utils/messaging.py @@ -282,7 +282,11 @@ def update(self, message: Message) -> Message | list[Message]: message_dict = {} for key, value in values.items(): if isinstance(value, dict) and 'value' in value: - message_dict[key] = value + # Shallow-copy to avoid mutating transformer's internal + # state (latest_input_struct) during trace injection. + message_dict[key] = {**value} + if 'metadata' in value: + message_dict[key]['metadata'] = {**value['metadata']} else: message_dict[key] = {'value': value} # Re-attach input metadata if this key had it @@ -292,11 +296,23 @@ def update(self, message: Message) -> Message | list[Message]: message_dict[key]['metadata'] = existing_meta self.transformer.updated = False + + # Aggregate trace_ids from ALL contributing inputs, + # not just the triggering message. + parent_ids = {message.trace_id} + input_structs = getattr(self.transformer, 'latest_input_struct', {}) + for struct in (input_structs or {}).values(): + if isinstance(struct, dict): + trace_info = (struct.get('metadata') or {}).get('trace') or {} + tid = trace_info.get('trace_id') + if tid: + parent_ids.add(tid) + out_msg = Message( topic=self.topic, source=str(self), value=message_dict, - parent_trace_ids=[message.trace_id], + parent_trace_ids=list(parent_ids), ) # Inject trace_id into each variable struct metadata for key in out_msg.value: @@ -587,11 +603,20 @@ def update(self, message: Message) -> list[Message]: # logger.debug(f"not packing output passign raw: {pred}") output = pred + # Aggregate trace_ids from ALL input variables, not just + # the triggering message. + parent_ids = {message.trace_id} + for meta in input_meta.values(): + trace_info = (meta.get('trace') or {}) + tid = trace_info.get('trace_id') + if tid: + parent_ids.add(tid) + out_msg = Message( topic=self.topic, source=str(self), value=output, - parent_trace_ids=[message.trace_id], + parent_trace_ids=list(parent_ids), ) # Inject trace_id and merge input metadata into each output variable struct for key in out_msg.value: diff --git a/poly_lithic/src/utils/readme.md b/poly_lithic/src/utils/readme.md new file mode 100644 index 0000000..10f7f75 --- /dev/null +++ b/poly_lithic/src/utils/readme.md @@ -0,0 +1,1101 @@ + + +![Poly-Lithic Logo](./images/logo.svg) + +[![Main build and test workflow](https://github.com/ISISNeutronMuon/poly-lithic/actions/workflows/main.yml/badge.svg?branch=main)](https://github.com/ISISNeutronMuon/poly-lithic/actions/workflows/main.yml) + +[![Build & Push Docker Images on Tag](https://github.com/ISISNeutronMuon/poly-lithic/actions/workflows/docker-build.yml/badge.svg)](https://github.com/ISISNeutronMuon/poly-lithic/actions/workflows/docker-build.yml) + +[![Build & Deploy Sphinx Docs](https://github.com/ISISNeutronMuon/poly-lithic/actions/workflows/docs.yml/badge.svg)](https://github.com/ISISNeutronMuon/poly-lithic/actions/workflows/docs.yml) + + +# Table of Contents + - [Installation](#installation) + - [Usage](#usage) + - [Project Commands](#project-commands) + - [Initialise a new project](#initialise-a-new-project) + - [Initialise from an existing model](#initialise-from-an-existing-model) + - [Initialise from a JSON sample file](#initialise-from-a-json-sample-file) + - [Update an existing config](#update-an-existing-config) + - [Configuration file](#configuration-file-formerly-pv_mappings-files) + - [Plugin API](#plugin-api) + - [Modules](#modules) + - [Interface](#interface) + - [Interface Configs](#interface-configs) + - [p4p sample configuration](#p4p-sample-configuration) + - [p4p_server sample configuration](#p4p_server-sample-configuration) + - [k2eg sample configuration](#k2eg-sample-configuration) + - [fastapi_server sample configuration](#fastapi_server-sample-configuration) + - [REST API Endpoints](#rest-api-endpoints) + - [Job Lifecycle & Tracking](#job-lifecycle--tracking) + - [Transformer](#transformer) + - [Transformer Configs](#transformer-configs) + - [SimpleTransformer sample configuration](#simpletransformer-sample-configuration) + - [CAImageTransformer sample configuration](#caimagetransformer-sample-configuration) + - [CompoundTransformer sample configuration](#compoundtransformer-sample-configuration) + - [PassThroughTransformer sample configuration](#passthroughtransformer-sample-configuration) + - [Model](#model) + - [Model Configs](#model-configs) + - [Example Model - Local Model](#example-model---local-model) + - [Example Model - MLFLow Model](#example-model---mlflow-model) + + - [Roadmap](#roadmap) +# Poly-Lithic + +Poly-Lithic is a package that allows you do deploy any model with an arbitrary number of inputs and outputs, related data transformations and system interfaces. + +Each deployment is defined by a model, typically hosted and retrieved from [MLFlow](https://mlflow.org/) and YAML file that describes the DG (Directed Graph) of model, transformations and interfaces. There are no restrictions on the numbers and types of nodes in the graph, so it may be used for things other than ML models. + + +## Installation + +Python `3.11.x` recommended. + +```bash +pip install poly-lithic +``` +for development: + +```bash +pip install -r reqirements.txt +pip install -e . +``` + +Alternatively with uv: + +```bash +uv pip install poly-lithic +``` + +for development with uv: +```bash +uv pip install -r requirements.txt +uv pip install -e . +``` + +with docker: + +```bash +docker compose -f ./docker/docker_compose.yml up +``` + +## Usage + +```python +model_manager run --publish -c ./tests/pv_mapping_mlflow.yaml -e ./tests/env.json +``` +or +```python +pl run --publish -c ./tests/pv_mapping_mlflow.yaml -e ./tests/env.json +``` +The env file is a json file that contains the environment variables that are used in the deployment. In this example we are pulling the [torch](./examples/torch_and_generic.ipynb) model and wrapping it with simple transformers and a simple p4p server. + +Reqired variables are: +- `AWS_ACCESS_KEY_ID` +- `AWS_SECRET` +- `AWS_DEFAULT_REGION` +- `AWS_REGION` +- `MLFLOW_S3_ENDPOINT_URL` +- `MINIO_ROOT_PASSWORD` +- `MINIO_ROOT_USER` +- `MINIO_SITE_REGION` +- `MLFLOW_TRACKING_URI` +- `PUBLISH` - set to `true` for the deployment to publish data to the interface. This flag serves as a safety measure to prevent accidental publishing of data to live system. + +See [this](https://mlflow.org/docs/latest/api_reference/python_api/mlflow.environment_variables.html) for explantions of the MLFlow environment variables. + +## Project Commands + +The `project` command group scaffolds new deployment projects and updates existing configurations. It has two subcommands: `init` and `update`. + +``` +pl project --help +pl project init --help +pl project update --help +``` + +### Initialise a new project + +Create a blank deployment project with placeholder variables: + +```bash +pl project init --name my-model --interface p4p_server --model-source local +``` + +Use `--docker` and `--kubernetes` to include Docker and K8s manifests: + +```bash +pl project init -n my-model -i fastapi -m mlflow --docker --kubernetes +``` + +For non-interactive usage (CI, scripting), pass `--no-prompt`: + +```bash +pl project init --name my-model --no-prompt +``` + +### Initialise from an existing model + +When you already have a `model_definition.py`, pass `--model-file` to introspect it at generation time. The generator extracts input/output variable names, types, defaults, and ranges, and pre-populates the deployment config: + +```bash +pl project init --name lume-demo -f model_definition.py -i p4p_server --no-prompt +``` + +The introspector resolves variables in this order: + +1. **Module-level variables** — If the file defines `input_variables` and `output_variables` as top-level lists (of dicts or lume-base objects), those are used directly. +2. **Factory class fallback** — Otherwise, a factory class (default `ModelFactory`) is instantiated and `get_model()` is called to obtain the variables from the model instance. + +You can override the factory class name: + +```bash +pl project init --name my-proj -f model_def.py --factory-class MyFactory --no-prompt +``` + +#### Plain dict model files (no lume-base required) + +You can define variables as plain Python dicts — no external dependencies needed: + +```python +input_variables = [ + {"name": "x1", "type": "scalar", "default_value": 0.0, "value_range": [-1, 1]}, + {"name": "x2", "type": "scalar", "default_value": 0.0}, + {"name": "signal", "type": "waveform", "length": 256}, +] + +output_variables = [ + {"name": "y", "type": "scalar"}, + {"name": "image_out", "type": "image", "image_size": {"x": 64, "y": 48}}, +] + +class ModelFactory: + def get_model(self): + # your model class here + ... +``` + +Each dict must have a `"name"` key. Other supported keys: + +| Key | Required | Description | +|---|---|---| +| `name` | yes | Variable name | +| `type` | no | One of `scalar`, `waveform`, `array`, `image` (defaults to `scalar`) | +| `default_value` | no | Default value for the variable | +| `value_range` | no | `[min, max]` range | +| `length` | no | Length for `waveform`/`array` types | +| `image_size` | no | `{"x": width, "y": height}` for `image` types | + +#### lume-base compatibility + +lume-base variable objects are also fully supported. The variable type is inferred automatically from the class name: + +| lume-base class | Inferred type | +|---|---| +| `ScalarVariable`, `ScalarInputVariable`, `ScalarOutputVariable` | `scalar` | +| `ArrayVariable`, `ArrayInputVariable`, `ArrayOutputVariable` | `waveform` | +| `ImageVariable`, `ImageInputVariable`, `ImageOutputVariable` | `image` | + +A minimal lume-base example: + +```python +from lume_model.base import LUMEBaseModel +from lume_model.variables import ScalarVariable + +class MyModel(LUMEBaseModel): + def _evaluate(self, input_dict): + return {"y": input_dict["x1"] + input_dict["x2"]} + +class ModelFactory: + def __init__(self): + self.model = MyModel( + input_variables=[ + ScalarVariable(name="x1", default_value=0, value_range=[-1, 1]), + ScalarVariable(name="x2", default_value=0), + ], + output_variables=[ScalarVariable(name="y")], + ) + + def get_model(self): + return self.model +``` + +Install the lume extras to pull in the required dependencies: + +```bash +pip install poly-lithic[lume] # lume-model ≥ 2.0.0 +# or with torch support: +pip install poly-lithic[torch] # lume-model ≥ 2.0.0 + torch ≥ 2.6.0 +``` + +### Initialise from a JSON sample file + +If you don't have a model file yet but know what your inputs and outputs look like, you can provide a **JSON sample file** with `--sample-file`. Variable names and types are inferred automatically from the sample data: + +```bash +pl project init --name my-model -s sample.json -i fastapi --no-prompt +``` + +The JSON file must contain `"input"` and `"output"` keys. Each can be either a **named dict** or an **unnamed list**. + +#### Named dict format (recommended) + +Keys become variable names, values are used for type inference: + +```json +{ + "input": { + "x1": 1.0, + "x2": 2, + "signal": [0.1, 0.2, 0.3, 0.4], + "image": [[1, 2], [3, 4], [5, 6]] + }, + "output": { + "y": 0.0, + "spectrum": [0.0, 0.0, 0.0] + } +} +``` + +#### Unnamed list format + +Variables are named `input_0`, `input_1`, ..., `output_0`, etc.: + +```json +{ + "input": [1.0, 2.0, 3.0], + "output": [0.0] +} +``` + +#### Type inference rules + +| Sample value | Inferred type | Extra fields | +|---|---|---| +| `int` or `float` | `scalar` | `default_value` set to the sample value | +| `str` or `bool` | `scalar` | `default_value` set to the sample value | +| 1-D list (e.g. `[1, 2, 3]`) | `waveform` | `length`, `default_value` | +| 2-D list (e.g. `[[1, 2], [3, 4]]`) | `image` | `image_size: {x, y}`, `default_value` | +| Empty list `[]` | `scalar` | — | + +### Update an existing config + +If you already have a generated project with placeholder variable names, you can patch the `deployment_config.yaml` in-place using a model file or a sample file: + +```bash +pl project update deployment_config.yaml -f model_definition.py +pl project update deployment_config.yaml -s sample.json +``` + +A coloured unified diff of the changes is printed after each update so you can see exactly what was modified. If the config is already up to date, a "No changes made" message is shown instead. + +This introspects the model (or infers from sample JSON) and replaces: +- Interface variable entries (PV names / FastAPI variable definitions) — including type, length, and image_size +- Input transformer symbols and variable mappings +- Model output variables with their types +- Output transformer symbols and variable mappings + +Options `--model-file` and `--sample-file` are mutually exclusive. An optional `--factory-class` flag is available when using `--model-file` if your factory is not named `ModelFactory`. + +## Configuration file (formerly pv_mappings files) + +The configuration file consists of 2 sections `deployment` and `modules`. Former describes deployment type and other setings such as refresh rate. The latter describes the nodes the modules and their connections to each other. + +Poly-Lithic supports two deployment modes: **continuous** (polling) and **event-driven** (reactive). The diagrams below show how each mode works and which settings affect the flow. + +
+Continuous Mode Flow + +In continuous mode the main loop polls all input PVs at a fixed interval defined by `rate`. Every tick reads **all** inputs, pushes them through the transformer → model → output pipeline, and writes results back. + +```yaml +deployment: + type: "continuous" + rate: 1 # seconds between polling cycles + trace_buffer_size: 10000 # circular trace buffer depth + trace_port: 8100 # REST API port for trace queries +``` + +```mermaid +flowchart TD + START([Start]) --> BUILD[Build observers & broker\nfrom YAML config] + BUILD --> TRACE[Start trace server\non trace_port] + TRACE --> TIMER_INIT["last_read = now()"] + + TIMER_INIT --> CHECK_TIMER{"now() − last_read\n> rate?"} + + CHECK_TIMER -- Yes --> RESET["last_read = now()"] + RESET --> GET_ALL["broker.get_all()\nRead ALL input PVs\nvia interface.get_many()"] + GET_ALL --> ENQUEUE[Messages enqueued\nwith trace_id & timestamp] + + CHECK_TIMER -- No --> CHECK_QUEUE{"broker.queue\nnon-empty?"} + + ENQUEUE --> PARSE[broker.parse_queue] + CHECK_QUEUE -- Yes --> PARSE + + PARSE --> TRANSFORM["Input Transformer\nhandler() stores each PV value"] + TRANSFORM --> ALL_GATE{"All input\nsymbols\npresent?"} + ALL_GATE -- No --> SLEEP + ALL_GATE -- Yes --> EVAL_FORMULA["transform():\nevaluate formulas\ne.g. x = PV_A * 2 + 10"] + EVAL_FORMULA --> MODEL["Model.evaluate()\nunpack inputs → run inference\npack outputs"] + MODEL --> OUT_TRANSFORM["Output Transformer\nmap model outputs → PV names"] + OUT_TRANSFORM --> PUBLISH_CHECK{"PUBLISH\n== True?"} + PUBLISH_CHECK -- Yes --> WRITE["interface.put_many()\nwrite to output PVs"] + PUBLISH_CHECK -- No --> DISCARD[Discard output] + WRITE --> ONE_SHOT{"one_shot\nflag?"} + DISCARD --> ONE_SHOT + ONE_SHOT -- Yes --> EXIT([Exit]) + ONE_SHOT -- No --> SLEEP["await asyncio.sleep(0.01)"] + + CHECK_QUEUE -- No --> SLEEP + SLEEP --> CHECK_TIMER + + style GET_ALL fill:#2374ab,color:#fff + style EVAL_FORMULA fill:#4a9c6d,color:#fff + style MODEL fill:#d4773b,color:#fff + style WRITE fill:#7b4f9e,color:#fff +``` + +| Setting | Type | Default | Effect | +|---------|------|---------|--------| +| `rate` | float | — | Seconds between `get_all()` polls. Lower = more frequent reads. | +| `trace_buffer_size` | int | 10000 | Max messages kept in the circular trace buffer. | +| `trace_port` | int | 8100 | Port for the trace REST API server. | +| `--publish` / `-p` | flag | off | When set, output PVs are written. Otherwise results are discarded. | +| `--one-shot` / `-o` | flag | off | Exit after the first complete pipeline cycle (useful for debugging). | + +
+ +
+Event-Driven Mode Flow + +In event-driven mode PV monitors fire callbacks whenever an external client writes a new value. A seeding step reads all current values first so the transformer's "all inputs present" gate is satisfied from the first event. + +```yaml +deployment: + type: "event_driven" + min_monitor_interval: 0.01 # per-PV throttle in seconds + on_change_only: false # skip if value unchanged + trace_buffer_size: 10000 + trace_port: 8100 +``` + +```mermaid +flowchart TD + START([Start]) --> BUILD[Build observers & broker\nfrom YAML config] + BUILD --> TRACE[Start trace server\non trace_port] + TRACE --> SEED["Seed: broker.get_all()\nRead current PV values"] + SEED --> DRAIN["Drain queue:\nbroker.parse_queue()\nuntil empty"] + DRAIN --> SEED_NOTE["All transformer inputs\nnow initialised"] + SEED_NOTE --> REG["Register PV monitors\non each input interface"] + + REG --> WAIT["Main loop:\nawait asyncio.sleep(0.01)"] + WAIT --> POLL_Q{"broker.queue\nnon-empty?"} + POLL_Q -- No --> WAIT + POLL_Q -- Yes --> PARSE[broker.parse_queue] + + subgraph MONITOR_CB ["PV Monitor Callback (per PV)"] + direction TB + EXT["External write:\npvput PV_NAME value"] --> HANDLER["p4p Handler.put()\nfires _monitor_callbacks"] + HANDLER --> THROTTLE{"now() − last_fire\n< min_monitor_interval?"} + THROTTLE -- Yes --> SKIP_T[Skip — throttled] + THROTTLE -- No --> DEDUP{"on_change_only\nand value == last_value?"} + DEDUP -- Yes --> SKIP_D[Skip — duplicate] + DEDUP -- No --> FIRE["Create Message\ntopic='in_interface'\nvalue = {PV: {value: X}}"] + FIRE --> BROKER_Q["Append to\nbroker.queue"] + end + + BROKER_Q -.-> POLL_Q + + PARSE --> TRANSFORM["Input Transformer\nhandler() updates PV value"] + TRANSFORM --> ALL_GATE{"All input\nsymbols\npresent?\n(seeding guarantees yes)"} + ALL_GATE -- No --> WAIT + ALL_GATE -- Yes --> EVAL_FORMULA["transform():\nevaluate formulas"] + EVAL_FORMULA --> MODEL["Model.evaluate()\nunpack → infer → pack"] + MODEL --> OUT_TRANSFORM["Output Transformer\nmap outputs → PV names"] + OUT_TRANSFORM --> PUBLISH_CHECK{"PUBLISH\n== True?"} + PUBLISH_CHECK -- Yes --> WRITE["interface.put_many()\nwrite to output PVs"] + PUBLISH_CHECK -- No --> DISCARD[Discard output] + WRITE --> ONE_SHOT{"one_shot\nflag?"} + DISCARD --> ONE_SHOT + ONE_SHOT -- Yes --> EXIT([Exit]) + ONE_SHOT -- No --> WAIT + + style SEED fill:#2374ab,color:#fff + style EXT fill:#c94c4c,color:#fff + style FIRE fill:#2374ab,color:#fff + style EVAL_FORMULA fill:#4a9c6d,color:#fff + style MODEL fill:#d4773b,color:#fff + style WRITE fill:#7b4f9e,color:#fff +``` + +| Setting | Type | Default | Effect | +|---------|------|---------|--------| +| `min_monitor_interval` | float | 0.0 | Per-PV throttle — callbacks arriving faster than this interval are skipped. | +| `on_change_only` | bool | false | When true, a callback is skipped if the new value equals the previous one. | +| `trace_buffer_size` | int | 10000 | Max messages kept in the circular trace buffer. | +| `trace_port` | int | 8100 | Port for the trace REST API server. | +| `--publish` / `-p` | flag | off | When set, output PVs are written. Otherwise results are discarded. | +| `--one-shot` / `-o` | flag | off | Exit after the first complete pipeline cycle. | + +
+ + +#### Example configuration file +```yaml +deployment: + type: "continuous" #type of deployment, continuous is the only one supported at the moment but more will be added + rate: 0.25 #refresh rate in seconds + +modules: + module1: + name: "module1" # name of the module used to identify it in the graph + type: "type.subtype" # type of the module, used to identify the module class and subclass + pub: "topic1" # topic the outputs will be published to, similar to MQTT, Kafka, ROS etc + sub: # topics the module will subscribe to, we listen for and transform data from these topics + - "update" # update is a special topic that will trigger an interface module to run get_all method (get_many for all keys) + - "topic3" + module_args: None # defines what arguments to pass to the module observer, if any this can inform unpacking etc + config: # configuration specific to the module type + key1: "value1" + keyn: "valuen" + + module2: + ... + pub: "topic2" + sub: + - "topic1" + module3: + ... + pub: "topic3" + sub: + - "topic2" +``` +The graph for the above configuration would look like this: + +```mermaid +graph TD; + every_0.25s --> module1 + module1 --> module2 + module2 --> module3 + module3 --> module1 +``` + +Under the hood we are passing messages in the format: + +```json +{ + "topic": "topic1", + "data": { + "key1": {"value" : 1}, + "key2": {"value" : [1,2,3]}, + "key3": {"value" : {...}} + } +} +``` +Note that the data is a dictionary of dictionaries. + +## Plugin API + +Read more in the [plugin readme](./poly_lithic/templates/plugin/README.md). + + +## Modules + +### Interface +Interface modules are used to interact with external data, usually an accelerators control systems but can be anything. They follow the following structure (see [base interface class](./poly-lithic/src/interfaces/BaseInterface.py)): + +```python +class BaseInterface(ABC): + @abstractmethod + def __init__(self, config): + pass + + @abstractmethod + def monitor(self, name, handler, **kwargs): # not used at the moment but will be used to monitor the interface for changes, rather than polling when p4p can monitor more than 4 pv's + pass + + @abstractmethod + def get(self, name, **kwargs): # get a value from the interface + pass + + @abstractmethod + def put(self, name, value, **kwargs): # put a value to the interface + pass + + @abstractmethod + def put_many(self, data, **kwargs): # put many values to the interface + pass + + @abstractmethod + def get_many(self, data, **kwargs): # get many values from the interface + pass +``` +All values are expected to come in as dictionaries of dictionaries with the following format: + +```python +# for sigular puts and gets +name = "key1" +value = {"value" : 1, "timestamp": 1234567890, "metadata": "some meta data"} # note tha the timestamp and metadata are optional and unusued at the moment + +# for _many +data = { + "key1": {"value" : 1, "timestamp": 1234567890, "metadata": "some meta data"}, + "key2": {"value" : [1,2,3]}, + "key3": {"value" : {...}} +} +``` +#### Interface Configs +| Module | Description | YAML configuration | +| ------ | ----------- | ------------------ | +| `p4p` | EPICS data source, must have an external EPICS server running. Note that SoftIOCPVA will not work with this module. | [config](#p4p-sample-configuration) | +| `p4p_server` | EPICS data source, host EPICS p4p server for specifed PVs | same [config](#p4p-sample-configuration) as `p4p`| +| `k2eg` | Kafka to EPICS gateway, get data from Kafka and write it to EPICS | [config](#k2eg-sample-configuration) | +| `fastapi_server` | HTTP/REST interface with job queue for request-response model inference | [config](#fastapi_server-sample-configuration) | + +##### `p4p` sample configuration +```yaml +deployment: + ... +modules: + mymodule: + ... + config: + EPICS_PVA_NAME_SERVERS: "epics.server.co.uk:5075" + # other EPICS_CONFIGS can go here + variables: + MY_VAR:TEST_A: + proto: pva + name: MY_VAR:TEST_A # name here is redundant as the name is the key in the variables dictionary, it will be removed in future versions + MY_VAR:TEST_B: + proto: pva + name: MY_VAR:TEST_B + MY_VAR:TEST_S: + proto: pva + name: MY_VAR:TEST_S + # default: 0 | [0.0, ... ,0.0] | no defaults for images optional + # type: scalar | waverform | image (default scalar) optional + # compute_alarm: true|false (default false) optional + # display/control/valueAlarm: native NTScalar metadata optional +``` + +##### `p4p_server` sample configuration +```yaml + config: + EPICS_PVA_NAME_SERVERS: "epics.server.co.uk:5075" + # other EPICS_CONFIGS can go here + variables: + MY_VAR:TEST_A: + proto: pva + name: MY_VAR:TEST_A + MY_VAR:TEST_B: + proto: pva + name: MY_VAR:TEST_B + MY_VAR:TEST_S: + proto: pva + name: MY_VAR:TEST_S + # default: 0 | [0.0, ... ,0.0] | no defaults for images optional + # type: scalar | waverform | image (default scalar) optional + # compute_alarm: true|false (default false) optional + # display/control/valueAlarm: native NTScalar metadata optional +``` +Yes, it is identical to p4p, the only difference is that the p4p server will host the PVs for the specified variables. + +###### Alarm config for `p4p` / `p4p_server` + +Scalar PVs can compute EPICS alarm fields from ``valueAlarm`` limits: + +- ``compute_alarm`` (bool, default ``false``) +- ``display`` (optional): ``limitLow``, ``limitHigh``, ``description``, ``format``, ``units`` +- ``control`` (optional): ``limitLow``, ``limitHigh``, ``minStep`` +- ``valueAlarm`` (optional native NT block) + +When ``compute_alarm: true``: + +- ``valueAlarm.active`` defaults to ``true`` if omitted +- ``valueAlarm.active: false`` is rejected +- required limits: ``lowAlarmLimit``, ``lowWarningLimit``, ``highWarningLimit``, ``highAlarmLimit`` +- optional severities default to: + ``lowAlarmSeverity=2``, ``lowWarningSeverity=1``, ``highWarningSeverity=1``, ``highAlarmSeverity=2`` + +Status mapping follows EPICS ``menuAlarmStat``: +``NO_ALARM=0``, ``HIHI=3``, ``HIGH=4``, ``LOLO=5``, ``LOW=6``. + +Notes: + +- Non-scalar PVs do not compute alarms. +- Non-scalar and scalar PVs may still pass explicit ``alarm`` payloads manually. +- Explicit ``alarm`` payload always overrides computed alarm. +- ``p4p`` client attempts structured put first; if the server rejects it, it retries with value-only put. + +###### Alarm evaluation flow + +```mermaid +flowchart TD + A[Incoming put payload] --> B{Explicit alarm in payload?} + B -- Yes --> C[Forward payload as-is] + B -- No --> D{Scalar PV?} + D -- No --> E[No computation
Write value only] + D -- Yes --> F{compute_alarm true?} + F -- No --> E + F -- Yes --> G{valueAlarm configured
active not false
and limits valid?} + G -- No --> H[Config validation error at startup] + G -- Yes --> I[Evaluate thresholds] + I --> J[Set alarm severity/status/message] + J --> K[Write structured payload] + C --> L{Client put accepted?} + K --> L + L -- Yes --> M[Done] + L -- No --> N[Client fallback:
retry value-only put] + N --> M +``` + +###### Model-side alarm override + +The model may override alarm fields by returning structured output: + +```python +return { + "ML:LOCAL:TEST_S": { + "value": output_value, + "alarm": {"severity": 2, "status": 3, "message": "HIHI (model override)"}, + } +} +``` + +This is supported by ``ModelObserver`` and passed through to interfaces. +In ``examples/base/local/deployment_config_p4p_alarm.yaml`` this now goes through an +``output_transformer`` direct-symbol mapping (``ML:LOCAL:TEST_S -> ML:LOCAL:TEST_S``), +which preserves ``alarm`` and other non-``value`` fields. +See runnable example: + +- config: ``examples/base/local/deployment_config_p4p_alarm.yaml`` +- model: ``examples/base/local/model_definition_alarm_override.py`` + +#### `k2eg` Sample configuration + +This module is built on top of SLAC's [k2eg](https://github.com/slaclab/k2eg), it's great because it allows you get data from `pva` and `ca` protocols over Kafka. currently its the only interface that supports `ca` protocol. + +```yaml +input_data: + get_method: "k2eg" + config: + variables: + MY_VAR:TEST_A: + proto: ca # supports ca or pva + name: MY_VAR:TEST_A + MY_VAR:TEST_B: + proto: pva + name: MY_VAR:TEST_B +``` + +##### `fastapi_server` sample configuration + +The `fastapi_server` interface exposes a REST API for submitting inference jobs and retrieving results. It manages an internal job queue and variable store, and embeds a uvicorn server. + +> Warning: `fastapi_server` is experimental and may change or be removed without notice. + +###### Config fields + +| Field | Type | Default | Description | +| ----- | ---- | ------- | ----------- | +| `name` | string | `"fastapi_server"` | Display name | +| `host` | string | `"127.0.0.1"` | Bind address | +| `port` | int | `8000` | Bind port | +| `start_server` | bool | `true` | Whether to launch embedded uvicorn | +| `wait_for_server_start` | bool | `false` | Block until server is accepting connections | +| `startup_timeout_s` | float | `2.0` | Max wait for startup | +| `input_queue_max` | int | `1000` | Max queued jobs before rejecting (HTTP 429) | +| `output_queue_max` | int | `1000` | Max completed jobs before oldest is evicted | +| `cors_origins` | list\[string\] | `[]` | CORS allow-origins (empty = no CORS middleware) | +| `variables` | dict | **required** | Variable definitions (see below) | + +###### Variable fields + +| Field | Type | Default | Description | +| ----- | ---- | ------- | ----------- | +| `mode` | string | `"inout"` | `in`, `out`, or `inout` | +| `type` | string | `"scalar"` | `scalar`, `waveform`, `array`, or `image` | +| `default` | any | `0.0` / zeros | Initial value (not supported for `image` type) | +| `length` | int | `10` | Array/waveform length when no default is provided | +| `image_size` | dict | — | Required for `image` type: `{"x": int, "y": int}` | + +###### Example YAML + +```yaml +modules: + my_fastapi: + name: "my_fastapi" + type: "interface.fastapi_server" + pub: "in_interface" + sub: + - "get_all" + - "out_transformer" + config: + name: "my_fastapi_interface" + host: "127.0.0.1" + port: 8000 + start_server: true + input_queue_max: 1000 + output_queue_max: 1000 + cors_origins: + - "http://localhost:3000" + variables: + MY_INPUT_A: + mode: in + type: scalar + default: 0.0 + MY_INPUT_B: + mode: in + type: array + default: [1, 2, 3, 4, 5] + MY_IMAGE_IN: + mode: in + type: image + image_size: + x: 128 + y: 64 + MY_OUTPUT: + mode: out + type: scalar + default: 0.0 +``` +###### Runnable config (array/waveform) + +For a full runnable config that includes `array` and `waveform` variables, see +[examples/base/local/deployment_config_fastapi_array_waveform.yaml](examples/base/local/deployment_config_fastapi_array_waveform.yaml). + +Run it with: + +```bash +pl run --publish -c examples/base/local/deployment_config_fastapi_array_waveform.yaml +``` + +Sample curl commands for the runnable config: + +```bash +curl http://127.0.0.1:8000/health + +curl -X POST http://127.0.0.1:8000/submit \ + -H 'Content-Type: application/json' \ + -d '{"job_id":"job-001","variables":{"ML:LOCAL:WAVEFORM_IN":{"value":[1,2,3,4,5,6,7,8]},"ML:LOCAL:ARRAY_IN":{"value":[10,11,12,13]},"ML:LOCAL:TEST_A":{"value":1.23},"ML:LOCAL:TEST_B":{"value":4.56}}}' + +curl -X POST http://127.0.0.1:8000/get \ + -H 'Content-Type: application/json' \ + -d '{"variables":["ML:LOCAL:WAVEFORM_IN","ML:LOCAL:ARRAY_IN","ML:LOCAL:TEST_S","ML:LOCAL:WAVEFORM_OUT"]}' + +curl http://127.0.0.1:8000/jobs/next + +curl http://127.0.0.1:8000/jobs/job-001 +``` + +###### REST API Endpoints + +| Method | Path | Description | +| ------ | ---- | ----------- | +| `GET` | `/health` | Health check — returns `{"status": "ok", "type": "interface.fastapi_server"}` | +| `GET` | `/settings` | Variable metadata, queue limits, and route table | +| `POST` | `/submit` | Submit a single inference job | +| `POST` | `/get` | Read current variable values | +| `POST` | `/jobs` | Submit a batch of jobs | +| `GET` | `/jobs/next` | Dequeue the next completed job | +| `GET` | `/jobs/{job_id}` | Get the status of a specific job | + +**Submit request body:** +```json +{ + "job_id": "optional-custom-id", + "variables": { + "MY_INPUT_A": {"value": 3.14}, + "MY_INPUT_B": {"value": [10, 20, 30, 40, 50]} + } +} +``` + +**Job snapshot response (from `/jobs/next` or `/jobs/{job_id}`):** +```json +{ + "job_id": "uuid", + "status": "completed", + "submitted_at": 1707600000.0, + "started_at": 1707600001.0, + "completed_at": 1707600002.0, + "error": null, + "inputs": {"MY_INPUT_A": {"value": 3.14}}, + "outputs": {"MY_OUTPUT": {"value": 42.0}} +} +``` + +**Error codes:** + +| Code | Condition | +| ---- | --------- | +| 403 | Write to a read-only variable (`mode: out`) | +| 404 | Unknown variable name, unknown job ID, or no completed jobs for `/jobs/next` | +| 409 | Duplicate job ID | +| 422 | Type validation failure (e.g. wrong shape, non-numeric value) | +| 429 | Input queue full | + +###### Job Lifecycle & Tracking + +Jobs submitted via `/submit` or `/jobs` follow this lifecycle: + +``` +submit → queued → running → completed +``` + +1. **Queued** — the job is validated and placed in the input queue. +2. **Running** — on each clock tick, **one** queued job is transitioned to running and its input values are loaded into the variable store for the pipeline to process. +3. **Completed** — when the pipeline writes results back via `put_many`, the oldest running job is marked as completed and its outputs are recorded. + +Completed jobs can be retrieved via `GET /jobs/next` (FIFO dequeue) or `GET /jobs/{job_id}` (by ID). + +> [!NOTE] +> **Current tracking limitation (Stage 1 / v1.7.3+):** +> Job tracking is currently **approximated using FIFO ordering**. The pipeline's transformers strip message metadata, so the `job_id` is typically not propagated through to `put_many`. Instead, the system uses a FIFO fallback: when results arrive, the **oldest running job** is assumed to be the one that completed. To enforce this assumption, the clock-driven path transitions only **one queued job per tick** to running state. +> +> This approach is reliable for single-job-at-a-time workloads but does not support true concurrent job tracking. +> +> **Planned improvement (Stage 2 / v1.8+):** +> Proper job tracking will be integrated via trace propagation across the message broker. Each job's `job_id` will be carried through the full pipeline in struct metadata, enabling accurate matching of results to jobs even under concurrent load. + +### Transformer + +Transformers are used to transform data from one format to another, they can be used to perform some data processing, aggregation or any other transformation action. They follow the structure (see [base transformer class](./poly-lithic/src/transformers/BaseTransformer.py)): + +```python +class BaseTransformer: + @abstractmethod + def __init__(self, config: dict): + """ + config: dict passed from the pv_mappings.yaml files. + """ + pass + + @abstractmethod + def transform(self): + """ + Call transform function to transform the input data, see SimpleTransformer in model_manager/src/transformers/BaseTransformers.py for an example. + """ + pass + + @abstractmethod + def handler(self, pv_name: str, value: dict | float | int): + """ + Handler function to handle the input data, in most cases it initiates the transform function when all the input data is available. + Handler is the only function exposed to the main loop of the program aside from initial configuration. + """ + pass +``` +#### Transformer Configs +| Module | Description | YAML configuration | +| ------ | ----------- | ------------------ | +| `SimpleTransformer` | Simple transformer that can be used to transform scalar values (ca or pv values that have a `value` field) | [config](#simpletransformer-sample-configuration) | +| `CAImageTransformer` | Transformer that can be used to transform a triplet of an array, x and y ca values into a np array | [config](#caimagetransformer-sample-configuration) | +| `CompoundTransformer` | Compound transformer that can be used to have multiple transformers in parallel | [config](#compoundtransformer-sample-configuration) | +| `PassThroughTransformer` | Transformer that can be used to pass data through without any transformation other than the tag | [config](#passthroughtransformer-sample-configuration) | + +##### `SimpleTransformer` Sample configuration +```yaml +modules: + input_transformer: + name: "input_transformer" + type: "transformer.SimpleTransformer" + pub: "model_input" + sub: + - "system_input" + module_args: None + config: + symbols: + - "LUME:MLFLOW:TEST_B" + - "LUME:MLFLOW:TEST_A" + variables: + x2: + formula: "LUME:MLFLOW:TEST_B" + x1: + formula: "LUME:MLFLOW:TEST_A" +``` + +##### `CAImageTransformer` Sample configuration +```yaml +modules: + image_transformer: + name: "image_transformer" + type: "transformer.CAImageTransformer" + pub: "model_input" + sub: + - "update" + module_args: None + config: + variables: + img_1: + img_ch: "MY_TEST_CA" + img_x_ch: "MY_TEST_CA_X" + img_y_ch: "MY_TEST_CA_Y" + img_2: + img_ch: "MY_TEST_C2" + img_x_ch: "MY_TEST_CA_X2" + img_y_ch: "MY_TEST_CA_Y2" +``` + +##### `PassThroughTransformer` Sample configuration +```yaml +modules: + output_transformer: + name: "output_transformer" + type: "transformer.PassThroughTransformer" + pub: "system_output" + sub: + - "model_output" + module_args: None + config: + variables: + LUME:MLFLOW:TEST_IMAGE: "y_img" +``` +##### `CompoundTransformer` Sample configuration +> [!CAUTION] +> This module will be deprecated in the future, pub-sub model means that compound transformers are no longer needed. + +```yaml +modules: + compound_transformer: + name: "compound_transformer" + type: "transformer.CompoundTransformer" + pub: "model_input" + sub: + - "update" + module_args: None + config: + transformers: + transformer_1: + type: "SimpleTransformer" + config: + symbols: + - "MY_TEST_A" + - "MY_TEST_B" + variables: + x2: + formula: "MY_TEST_A*2" + x1: + formula: "MY_TEST_B+MY_TEST_A" + transformer_2: + type: "CAImageTransformer" + config: + variables: + img_1: + img_ch: "MY_TEST_CA" + img_x_ch: "MY_TEST_CA_X" + img_y_ch: "MY_TEST_CA_Y" + img_2: + img_ch: "MY_TEST_C2" + img_x_ch: "MY_TEST_CA_X2" + img_y_ch: "MY_TEST_CA_Y2" +``` + + +### Model + +Models are the core of the deployment, they can be retrieved locally or from MLFlow and accept data in form of dictionries. By deafault models pivot the dictionry or rather remove the additional keys from messages to simplify the data structure that the model has to process. + +All models have to implement the `evaluate` method that takes a dictionary of inputs and returns a dictionary of outputs. + +#### Model Config + +```yaml +model: # this is the name of the model module, it is used to identify the model in the graph + name: "model" # name of the model used to identify it in the graph, overrides the name in the module section + type: "model.SimpleModel" # type of module, used to identify the model class and subclass, in this case we are saying it a model + pub: "model" # where the model will publish its outputs, this is the topic that the model will publish to + sub: "in_transformer" # topic that the model will subscribe to, this is the topic that the model will listen for inputs + module_args: None # defines what arguments to pass to the model observer, if any this can inform unpacking etc + config: + type: "modelGetter" # defines the type of model getter, this is used to identify the model getter class + args: # arguments to pass to the model getter class, in this case we are passing the path to the model definition file + +``` +See the following examples for usage + + +###### Example Model - Local Model + +```python + +class SimpleModel(torch.nn.Module): + def __init__(self): + super(SimpleModel, self).__init__() + self.linear1 = torch.nn.Linear(2, 10) + self.linear2 = torch.nn.Linear(10, 1) + + def forward(self, x): # this is for our benefit, it is not used by poly-lithic + x = torch.relu(self.linear1(x)) + x = self.linear2(x) + return x + + # this method is necessary for the model to be evaluated by poly-lithic + def evaluate(self, x: dict) -> dict: + # x will be a dicrt of keys and values + # {"x": x, "y": y} + input_tensor = torch.tensor([x['x'], x['y']], dtype=torch.float32) + # you may want to do somethinf more complex here + output_tensor = self.forward(input_tensor) + # return a dictionary of keys and values + return {'output': output_tensor.item()} + ``` + + Lets say we want to retreive the model locally, we need to specify a factory class: + +```python + +class ModelFactory: + # can do more complex things here but we will just load the model from a locally saved file + def __init__(self): + # add this path to python environment + os.environ['PYTHONPATH'] = os.path.abspath( + os.path.join(os.path.dirname(__file__), '..', '..', '..') + ) + print('PYTHONPATH set to:', os.environ['PYTHONPATH']) + self.model = SimpleModel() + model_path = 'examples/base/local/model.pth' + if os.path.exists(model_path): + self.model.load_state_dict(torch.load(model_path)) + print('Model loaded successfully.') + else: + print( + f"Warning: Model file '{model_path}' not found. Using untrained model." + ) + print('ModelFactory initialized') + + # this method is necessary for the model to be retrieved by poly-lithic + def get_model(self): + return self.model +``` +The in the config file: + +```yaml +... +model: + name: "model" + type: "model.SimpleModel" + pub: "model" + sub: "in_transformer" + module_args: None + config: + type: "LocalModelGetter" + args: + model_path: "examples/base/local/model_definition.py" # path to the model definition + model_factory_class: "ModelFactory" # class that you use to create the model + variables: + max: + type: "scalar" +... +``` + +Then to run the model: +```bash +pl run --publish -c examples/base/local/deployment_config.yaml +``` +See the [local example notebook](./examples/base/simple_model_local.ipynb) for more details. + +#### Example Model - MLFLow Model +See the [MLFlow example notebook](./examples/base/simple_model_mlflow.ipynb) for more details. diff --git a/readme.md b/readme.md index 563a040..e1a9e40 100644 --- a/readme.md +++ b/readme.md @@ -302,6 +302,145 @@ Options `--model-file` and `--sample-file` are mutually exclusive. An optional ` The configuration file consists of 2 sections `deployment` and `modules`. Former describes deployment type and other setings such as refresh rate. The latter describes the nodes the modules and their connections to each other. +Poly-Lithic supports two deployment modes: **continuous** (polling) and **event-driven** (reactive). The diagrams below show how each mode works and which settings affect the flow. + +
+Continuous Mode Flow + +In continuous mode the main loop polls all input PVs at a fixed interval defined by `rate`. Every tick reads **all** inputs, pushes them through the transformer → model → output pipeline, and writes results back. + +```yaml +deployment: + type: "continuous" + rate: 1 # seconds between polling cycles + trace_buffer_size: 10000 # circular trace buffer depth + trace_port: 8100 # REST API port for trace queries +``` + +```mermaid +flowchart TD + START([Start]) --> BUILD[Build observers & broker\nfrom YAML config] + BUILD --> TRACE[Start trace server\non trace_port] + TRACE --> TIMER_INIT["last_read = now()"] + + TIMER_INIT --> CHECK_TIMER{"now() − last_read\n> rate?"} + + CHECK_TIMER -- Yes --> RESET["last_read = now()"] + RESET --> GET_ALL["broker.get_all()\nRead ALL input PVs\nvia interface.get_many()"] + GET_ALL --> ENQUEUE[Messages enqueued\nwith trace_id & timestamp] + + CHECK_TIMER -- No --> CHECK_QUEUE{"broker.queue\nnon-empty?"} + + ENQUEUE --> PARSE[broker.parse_queue] + CHECK_QUEUE -- Yes --> PARSE + + PARSE --> TRANSFORM["Input Transformer\nhandler() stores each PV value"] + TRANSFORM --> ALL_GATE{"All input\nsymbols\npresent?"} + ALL_GATE -- No --> SLEEP + ALL_GATE -- Yes --> EVAL_FORMULA["transform():\nevaluate formulas\ne.g. x = PV_A * 2 + 10"] + EVAL_FORMULA --> MODEL["Model.evaluate()\nunpack inputs → run inference\npack outputs"] + MODEL --> OUT_TRANSFORM["Output Transformer\nmap model outputs → PV names"] + OUT_TRANSFORM --> PUBLISH_CHECK{"PUBLISH\n== True?"} + PUBLISH_CHECK -- Yes --> WRITE["interface.put_many()\nwrite to output PVs"] + PUBLISH_CHECK -- No --> DISCARD[Discard output] + WRITE --> ONE_SHOT{"one_shot\nflag?"} + DISCARD --> ONE_SHOT + ONE_SHOT -- Yes --> EXIT([Exit]) + ONE_SHOT -- No --> SLEEP["await asyncio.sleep(0.01)"] + + CHECK_QUEUE -- No --> SLEEP + SLEEP --> CHECK_TIMER + + style GET_ALL fill:#2374ab,color:#fff + style EVAL_FORMULA fill:#4a9c6d,color:#fff + style MODEL fill:#d4773b,color:#fff + style WRITE fill:#7b4f9e,color:#fff +``` + +| Setting | Type | Default | Effect | +|---------|------|---------|--------| +| `rate` | float | — | Seconds between `get_all()` polls. Lower = more frequent reads. | +| `trace_buffer_size` | int | 10000 | Max messages kept in the circular trace buffer. | +| `trace_port` | int | 8100 | Port for the trace REST API server. | +| `--publish` / `-p` | flag | off | When set, output PVs are written. Otherwise results are discarded. | +| `--one-shot` / `-o` | flag | off | Exit after the first complete pipeline cycle (useful for debugging). | + +
+ +
+Event-Driven Mode Flow + +In event-driven mode PV monitors fire callbacks whenever an external client writes a new value. A seeding step reads all current values first so the transformer's "all inputs present" gate is satisfied from the first event. + +```yaml +deployment: + type: "event_driven" + min_monitor_interval: 0.01 # per-PV throttle in seconds + on_change_only: false # skip if value unchanged + trace_buffer_size: 10000 + trace_port: 8100 +``` + +```mermaid +flowchart TD + START([Start]) --> BUILD[Build observers & broker\nfrom YAML config] + BUILD --> TRACE[Start trace server\non trace_port] + TRACE --> SEED["Seed: broker.get_all()\nRead current PV values"] + SEED --> DRAIN["Drain queue:\nbroker.parse_queue()\nuntil empty"] + DRAIN --> SEED_NOTE["All transformer inputs\nnow initialised"] + SEED_NOTE --> REG["Register PV monitors\non each input interface"] + + REG --> WAIT["Main loop:\nawait asyncio.sleep(0.01)"] + WAIT --> POLL_Q{"broker.queue\nnon-empty?"} + POLL_Q -- No --> WAIT + POLL_Q -- Yes --> PARSE[broker.parse_queue] + + subgraph MONITOR_CB ["PV Monitor Callback (per PV)"] + direction TB + EXT["External write:\npvput PV_NAME value"] --> HANDLER["p4p Handler.put()\nfires _monitor_callbacks"] + HANDLER --> THROTTLE{"now() − last_fire\n< min_monitor_interval?"} + THROTTLE -- Yes --> SKIP_T[Skip — throttled] + THROTTLE -- No --> DEDUP{"on_change_only\nand value == last_value?"} + DEDUP -- Yes --> SKIP_D[Skip — duplicate] + DEDUP -- No --> FIRE["Create Message\ntopic='in_interface'\nvalue = {PV: {value: X}}"] + FIRE --> BROKER_Q["Append to\nbroker.queue"] + end + + BROKER_Q -.-> POLL_Q + + PARSE --> TRANSFORM["Input Transformer\nhandler() updates PV value"] + TRANSFORM --> ALL_GATE{"All input\nsymbols\npresent?\n(seeding guarantees yes)"} + ALL_GATE -- No --> WAIT + ALL_GATE -- Yes --> EVAL_FORMULA["transform():\nevaluate formulas"] + EVAL_FORMULA --> MODEL["Model.evaluate()\nunpack → infer → pack"] + MODEL --> OUT_TRANSFORM["Output Transformer\nmap outputs → PV names"] + OUT_TRANSFORM --> PUBLISH_CHECK{"PUBLISH\n== True?"} + PUBLISH_CHECK -- Yes --> WRITE["interface.put_many()\nwrite to output PVs"] + PUBLISH_CHECK -- No --> DISCARD[Discard output] + WRITE --> ONE_SHOT{"one_shot\nflag?"} + DISCARD --> ONE_SHOT + ONE_SHOT -- Yes --> EXIT([Exit]) + ONE_SHOT -- No --> WAIT + + style SEED fill:#2374ab,color:#fff + style EXT fill:#c94c4c,color:#fff + style FIRE fill:#2374ab,color:#fff + style EVAL_FORMULA fill:#4a9c6d,color:#fff + style MODEL fill:#d4773b,color:#fff + style WRITE fill:#7b4f9e,color:#fff +``` + +| Setting | Type | Default | Effect | +|---------|------|---------|--------| +| `min_monitor_interval` | float | 0.0 | Per-PV throttle — callbacks arriving faster than this interval are skipped. | +| `on_change_only` | bool | false | When true, a callback is skipped if the new value equals the previous one. | +| `trace_buffer_size` | int | 10000 | Max messages kept in the circular trace buffer. | +| `trace_port` | int | 8100 | Port for the trace REST API server. | +| `--publish` / `-p` | flag | off | When set, output PVs are written. Otherwise results are discarded. | +| `--one-shot` / `-o` | flag | off | Exit after the first complete pipeline cycle. | + +
+ #### Example configuration file ```yaml diff --git a/tests/test_trace_lineage.py b/tests/test_trace_lineage.py new file mode 100644 index 0000000..2cb4599 --- /dev/null +++ b/tests/test_trace_lineage.py @@ -0,0 +1,280 @@ +# Tests for trace lineage completeness. +# Verifies that parent_trace_ids aggregates trace_ids from ALL contributing +# inputs, not just the triggering message. + +import pytest + +from poly_lithic.src.utils.messaging import ( + Message, + TransformerObserver, + ModelObserver, +) +from poly_lithic.src.transformers.BaseTransformers import ( + SimpleTransformer, + PassThroughTransformer, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_value_struct(value, trace_id): + """Build a value struct with embedded trace metadata.""" + return { + 'value': value, + 'metadata': {'trace': {'trace_id': trace_id}}, + } + + +class _MockModel: + def evaluate(self, value): + return {'out': sum(value.values())} + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +@pytest.fixture +def simple_transformer(): + config = { + 'variables': { + 'x': {'formula': 'PV_A'}, + 'y': {'formula': 'PV_B'}, + }, + 'symbols': ['PV_A', 'PV_B'], + } + return SimpleTransformer(config) + + +@pytest.fixture +def passthrough_transformer(): + config = { + 'variables': { + 'OUT_A': 'PV_A', + 'OUT_B': 'PV_B', + }, + } + return PassThroughTransformer(config) + + +# --------------------------------------------------------------------------- +# TransformerObserver — SimpleTransformer +# --------------------------------------------------------------------------- + +class TestTransformerObserverLineage: + """TransformerObserver output should list trace_ids from ALL inputs.""" + + def test_seed_then_single_event(self, simple_transformer): + """Seed both PVs, then trigger via one PV change. + The output parent_trace_ids must contain both the seed trace_id + AND the event trace_id.""" + obs = TransformerObserver(simple_transformer, 'in_transformer') + + # --- Seed: simulate get_all delivering both PVs in one message --- + seed_msg = Message( + topic='in_interface', + source='test', + value={ + 'PV_A': _make_value_struct(1.0, 'seed-1'), + 'PV_B': _make_value_struct(2.0, 'seed-1'), + }, + ) + seed_out = obs.update(seed_msg) + assert seed_out is not None # transform fires (all inputs present) + + # --- Event: only PV_B changes --- + event_msg = Message( + topic='in_interface', + source='test', + value={ + 'PV_B': _make_value_struct(3.0, 'event-2'), + }, + ) + event_out = obs.update(event_msg) + assert event_out is not None + + pids = set(event_out.parent_trace_ids) + # Must contain the event message trace_id + assert event_msg.trace_id in pids + # Must contain PV_A's trace_id from seeding (still in latest_input_struct) + assert 'seed-1' in pids + + def test_both_pvs_updated_separately(self, simple_transformer): + """Two separate single-PV messages. Output should reference both.""" + obs = TransformerObserver(simple_transformer, 'in_transformer') + + msg_a = Message( + topic='in_interface', + source='test', + value={'PV_A': _make_value_struct(1.0, 'trace-a')}, + ) + out_a = obs.update(msg_a) + # Only one input present — transformer should NOT fire + assert out_a is None + + msg_b = Message( + topic='in_interface', + source='test', + value={'PV_B': _make_value_struct(2.0, 'trace-b')}, + ) + out_b = obs.update(msg_b) + assert out_b is not None + + pids = set(out_b.parent_trace_ids) + assert 'trace-a' in pids + assert 'trace-b' in pids + + def test_seed_message_trace_id_also_included(self, simple_transformer): + """The message-level trace_id (not just embedded metadata) should + always appear in parent_trace_ids.""" + obs = TransformerObserver(simple_transformer, 'in_transformer') + + seed_msg = Message( + topic='in_interface', + source='test', + value={ + 'PV_A': _make_value_struct(1.0, 'meta-seed'), + 'PV_B': _make_value_struct(2.0, 'meta-seed'), + }, + ) + out = obs.update(seed_msg) + assert out is not None + assert seed_msg.trace_id in out.parent_trace_ids + + +# --------------------------------------------------------------------------- +# TransformerObserver — PassThroughTransformer +# --------------------------------------------------------------------------- + +class TestPassThroughTransformerLineage: + """Same lineage behaviour for PassThroughTransformer.""" + + def test_seed_then_single_event(self, passthrough_transformer): + obs = TransformerObserver(passthrough_transformer, 'out_transformer') + + seed_msg = Message( + topic='model', + source='test', + value={ + 'PV_A': _make_value_struct(10, 'pt-seed'), + 'PV_B': _make_value_struct(20, 'pt-seed'), + }, + ) + seed_out = obs.update(seed_msg) + assert seed_out is not None + + event_msg = Message( + topic='model', + source='test', + value={'PV_B': _make_value_struct(30, 'pt-event')}, + ) + event_out = obs.update(event_msg) + assert event_out is not None + + pids = set(event_out.parent_trace_ids) + assert event_msg.trace_id in pids + assert 'pt-seed' in pids + + +# --------------------------------------------------------------------------- +# ModelObserver +# --------------------------------------------------------------------------- + +class TestModelObserverLineage: + """ModelObserver output should aggregate trace_ids from input metadata.""" + + def test_multiple_input_traces(self): + model_obs = ModelObserver(model=_MockModel(), topic='model') + + msg = Message( + topic='in_transformer', + source='test', + value={ + 'x': _make_value_struct(1.0, 'input-trace-x'), + 'y': _make_value_struct(2.0, 'input-trace-y'), + }, + ) + results = model_obs.update(msg) + assert len(results) == 1 + out = results[0] + + pids = set(out.parent_trace_ids) + assert msg.trace_id in pids + assert 'input-trace-x' in pids + assert 'input-trace-y' in pids + + def test_single_input_trace(self): + """When all inputs share the same trace_id, no duplicates.""" + model_obs = ModelObserver(model=_MockModel(), topic='model') + + msg = Message( + topic='in_transformer', + source='test', + value={ + 'x': _make_value_struct(1.0, 'same-trace'), + 'y': _make_value_struct(2.0, 'same-trace'), + }, + ) + results = model_obs.update(msg) + out = results[0] + + # 'same-trace' + msg.trace_id — at most 2 entries + assert 'same-trace' in out.parent_trace_ids + assert msg.trace_id in out.parent_trace_ids + + def test_no_metadata_no_crash(self): + """Inputs without metadata should not cause errors.""" + model_obs = ModelObserver(model=_MockModel(), topic='model') + + msg = Message( + topic='in_transformer', + source='test', + value={ + 'x': {'value': 1.0}, + 'y': {'value': 2.0}, + }, + ) + results = model_obs.update(msg) + out = results[0] + + # Only the message trace_id + assert out.parent_trace_ids == [msg.trace_id] + + +# --------------------------------------------------------------------------- +# Transformer without latest_input_struct (e.g. CAImageTransformer) +# --------------------------------------------------------------------------- + +class TestTransformerWithoutInputStruct: + """Transformers that lack latest_input_struct should still work + (fallback to triggering message trace_id only).""" + + def test_getattr_fallback(self): + from poly_lithic.src.transformers.BaseTransformer import BaseTransformer + + class MinimalTransformer(BaseTransformer): + def __init__(self): + self.updated = False + self.latest_transformed = {} + + def transform(self): + pass + + def handler(self, pv_name, value): + self.latest_transformed = {'out': value} + self.updated = True + + t = MinimalTransformer() + obs = TransformerObserver(t, 'test_topic') + + msg = Message( + topic='input', + source='test', + value={'PV_X': {'value': 42}}, + ) + out = obs.update(msg) + assert out is not None + # Only the message trace_id — no latest_input_struct to read from + assert msg.trace_id in out.parent_trace_ids From 5fa670c32da626266fd0dd9b5015084279a3c062 Mon Sep 17 00:00:00 2001 From: MatInGit Date: Tue, 7 Apr 2026 14:13:17 +0100 Subject: [PATCH 3/4] fixed testing for p4p since it was flaky --- tests/test_event_driven.py | 1 + tests/test_p4p.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/test_event_driven.py b/tests/test_event_driven.py index b5139a7..9d140a5 100644 --- a/tests/test_event_driven.py +++ b/tests/test_event_driven.py @@ -357,6 +357,7 @@ def test_chain_without_seeding_requires_all_pvs( # Step 5: Real p4p client write triggers Handler.put → monitor callback # --------------------------------------------------------------------------- +@pytest.mark.flaky_p4p(retries=3, backoff_max=2.0) class TestRealClientWrite: """Use a real p4p Context client to write to the server and verify that Handler.put() fires monitor callbacks end-to-end.""" diff --git a/tests/test_p4p.py b/tests/test_p4p.py index 411bdfe..b1342e8 100644 --- a/tests/test_p4p.py +++ b/tests/test_p4p.py @@ -33,6 +33,7 @@ def setup(): process.kill() +@pytest.mark.flaky_p4p(retries=3, backoff_max=2.0) def test_SimplePVAInterface_init(): config = { 'variables': { From 4254eee7c7af4ed06a280b6af4a2cd636bff5576 Mon Sep 17 00:00:00 2001 From: MatInGit Date: Tue, 7 Apr 2026 14:20:21 +0100 Subject: [PATCH 4/4] Add shared pytest configuration for flaky test retries --- tests/conftest.py | 56 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 tests/conftest.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1251f0a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,56 @@ +"""Shared pytest configuration and helpers for poly-lithic tests.""" + +import random +import time +import pytest + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "flaky_p4p(retries=3, backoff_max=2.0): " + "retry p4p/network tests with random backoff", + ) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + """Retry tests marked with @pytest.mark.flaky_p4p on call-phase failure.""" + outcome = yield + report = outcome.get_result() + + if report.when != "call" or not report.failed: + return + + marker = item.get_closest_marker("flaky_p4p") + if marker is None: + return + + retries = marker.kwargs.get("retries", 3) + backoff_max = marker.kwargs.get("backoff_max", 2.0) + + for attempt in range(2, retries + 1): + wait = random.uniform(0.1, backoff_max * attempt) + terminal = item.config.pluginmanager.get_plugin("terminalreporter") + if terminal is not None: + terminal.write_line( + f" \u21bb {item.nodeid} failed (attempt {attempt - 1}/{retries}), " + f"retrying in {wait:.2f}s \u2026", + yellow=True, + ) + time.sleep(wait) + + # Re-run the test's setup and call phases + item.setup() + try: + item.runtest() + except Exception: + continue # still failing, try again + + # Passed on retry — override the report to mark it as passed + report.outcome = "passed" + report.longrepr = None + report.wasxfail = None + return + + # All retries exhausted — leave the original failure report as-is