diff --git a/block_configs/tiled_processor_block.yaml b/block_configs/tiled_processor_block.yaml new file mode 100644 index 0000000..3b5eb91 --- /dev/null +++ b/block_configs/tiled_processor_block.yaml @@ -0,0 +1,27 @@ +blocks: + - name: XPS Tiled → LSE Processor + description: > + Watches the local Tiled server for new XPS scans written by Block 1 and + forwards each shot_mean frame as a RawFrameEvent over ZMQ to LSE. + + operator: + # Pass-through: RawFrameEvent → RawFrameEvent, no computation. + # Reuses the existing TiledRawFrameOperator from arroyosas. + class: arroyosas.tiled.tiled_poller.TiledRawFrameOperator + + listeners: + - class: tr_ap_xps.tiled_listener.tiled_new_scan_listener_factory + kwargs: + tiled_uri: "http://tiled:8000" + tiled_prefix: "beamlines/bl931/processed" + root_segments: + - xps_processed_images + poll_interval_sec: 2.0 + frame_poll_interval_sec: 3.0 + unchanged_cycles_threshold: 5 + seen_scans_file: "/data/tiled_poller/seen_scans.txt" + + publishers: + - class: tr_ap_xps.zmq_publisher.create_zmq_frame_publisher + kwargs: + zmq_address: "tcp://0.0.0.0:5000" \ No newline at end of file diff --git a/block_configs/timepix931_processor_block.yaml b/block_configs/timepix931_processor_block.yaml new file mode 100644 index 0000000..2ba1223 --- /dev/null +++ b/block_configs/timepix931_processor_block.yaml @@ -0,0 +1,31 @@ +blocks: + - name: XPS Timepix Processor + description: > + Real-time AP-XPS processor for Timepix data. + Listens for splash_timepix ZMQ messages, computes shot_mean via XPSOperator, + and publishes to two independent sinks: + 1. XPSWSResultPublisher — live websocket updates to the arroyoXPS frontend + 2. XPSTiledPublisher — persists shot_mean arrays to the local Tiled server + + operator: + class: tr_ap_xps.pipeline.xps_operator.build_xps_operator + kwargs: + build_heatmaps: false + otlp_endpoint: "http://jaeger:4317" + + listeners: + - class: tr_ap_xps.timepix.xps_timepix_listener_factory + kwargs: + zmq_connection_address: "tcp://host.docker.internal:5657" + + publishers: + - class: tr_ap_xps.websockets.xps_ws_publisher_factory + kwargs: + ws_url: "ws://0.0.0.0:8001/xps_operator" + + - class: tr_ap_xps.xps_tiled_publisher.xps_tiled_result_publisher_factory + kwargs: + tiled_uri: "http://tiled:8000" + tiled_prefix: "beamlines/bl931/processed" + root_segments: + - xps_processed_images \ No newline at end of file diff --git a/block_configs/timepix_processor_block.yaml b/block_configs/timepix_processor_block.yaml deleted file mode 100644 index 91e2bee..0000000 --- a/block_configs/timepix_processor_block.yaml +++ /dev/null @@ -1,19 +0,0 @@ -blocks: - - name: XPS Timepix Operator - description: Real-time AP-XPS processor for Timepix data - - operator: - class: tr_ap_xps.pipeline.xps_operator.build_xps_operator - kwargs: - build_heatmaps: false - otlp_endpoint: "http://jaeger:4317" - - listeners: - - class: tr_ap_xps.timepix.xps_timepix_listener_factory - kwargs: - zmq_connection_address: "tcp://host.docker.internal:5657" - - publishers: - - class: tr_ap_xps.websockets.xps_ws_publisher_factory - kwargs: - ws_url: "ws://0.0.0.0:8001/xps_operator" \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 1301568..994a37a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -24,7 +24,7 @@ services: timepix_processor: # profiles: # - dev - command: "/block_configs/timepix_processor_block.yaml" + command: "/block_configs/timepix931_processor_block.yaml" build: context: . dockerfile: Dockerfile_processor @@ -33,6 +33,27 @@ services: - "8001:8001" volumes: - ./block_configs:/block_configs:ro,Z + environment: + - RESULTS_TILED_API_KEY=${TILED_SINGLE_USER_API_KEY} + - LOGGING_LEVEL=DEBUG + networks: + mle_net: + + + tiled_processor: + command: "/block_configs/tiled_processor_block.yaml" + build: + context: . + dockerfile: Dockerfile_processor + restart: unless-stopped + ports: + - "5001:5001" + volumes: + - ./block_configs:/block_configs:ro,Z + - ./data/tiled_poller:/data/tiled_poller:Z + environment: + - RESULTS_TILED_API_KEY=${TILED_SINGLE_USER_API_KEY} + - LOGGING_LEVEL=DEBUG networks: mle_net: diff --git a/pyproject.toml b/pyproject.toml index 7f8212f..f69c737 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ classifiers = [ dependencies = [ "arroyopy @ git+https://github.com/als-computing/arroyopy.git@refs/heads/open_telemetry", + "arroyosas @ git+https://github.com/als-computing/arroyosas.git@upgrade_arroyopy", "astropy", "dynaconf", "python-dotenv", @@ -41,7 +42,7 @@ dev = [ "flake8", "fakeredis", "pre-commit", - "pytest-asyncio", + "pytest-asyncio==1.3.0", "pytest-mock", "tiled[server]", ] diff --git a/src/tr_ap_xps/tiled_listener.py b/src/tr_ap_xps/tiled_listener.py new file mode 100644 index 0000000..1e6df49 --- /dev/null +++ b/src/tr_ap_xps/tiled_listener.py @@ -0,0 +1,331 @@ +""" +Tiled-based listener for Block 2. + +Watches a Tiled container for new scans written by XPSTiledPublisher (Block 1) +and emits RawFrameEvents downstream so that ZMQFramePublisher can forward them +to LSE. + +Design mirrors TiledPollingRedisListener in arroyosas — minimal changes. +""" + +import asyncio +import logging +import os +import time +from datetime import datetime + +import pytz +from arroyopy.listener import Listener +from arroyopy.operator import Operator +from tiled.client import from_uri +from tiled.client.node import Container + +from arroyosas.schemas import RawFrameEvent, SASStart, SASStop, SerializableNumpyArrayModel +from tr_ap_xps.log_utils import setup_logger + +setup_logger( + logging.getLogger("tr_ap_xps"), log_level=os.getenv("LOGGING_LEVEL", "INFO") +) +logging.getLogger("httpx").setLevel(logging.WARNING) + +CALIFORNIA_TZ = pytz.timezone("US/Pacific") + + +def build_tiled_url( + tiled_uri: str, + tiled_prefix: str | None, + root_segments: list[str], + uuid: str, + shot_index: int, + height: int, + width: int, +) -> str: + """Construct the Tiled array slice URL for a single xps_averaged_heatmaps frame.""" + now = datetime.now(CALIFORNIA_TZ) + date_path = f"{now.year}/{now.month:02d}/{now.day:02d}" + + parts: list[str] = [] + if tiled_prefix: + parts.append(tiled_prefix) + parts.extend(root_segments) + parts.append(date_path) + parts.append(uuid) + parts.append(ARRAY_KEY) + + array_path = "/".join(parts) + slice_param = f"{shot_index}:{shot_index + 1},0:{height},0:{width}" + return f"{tiled_uri}/api/v1/array/full/{array_path}?slice={slice_param}" + +logger = logging.getLogger(__name__) + +ARRAY_KEY = "xps_averaged_heatmaps" + + +class TiledNewScanListener(Listener): + """Polls a Tiled container for new scan UUIDs written by XPSTiledPublisher. + + For each new scan it finds, it emits: + SASStart → one RawFrameEvent per shot_mean frame → SASStop + + The full container path is resolved lazily each poll cycle so startup + succeeds even before Block 1 has written any data to Tiled. + + Args: + operator: Downstream operator to forward messages to. + tiled_client: Root Tiled Container (top of the server). + tiled_uri: Base URI of the Tiled server. + tiled_prefix: Optional top-level path prefix (e.g. "beamlines/bl931/processed"). + root_segments: Path segments after the prefix (e.g. ["xps_processed_images"]). + poll_interval_sec: Seconds to sleep between checks for new scans. + frame_poll_interval_sec: Seconds to sleep between checks for new frames. + """ + + def __init__( + self, + operator: Operator, + tiled_uri: str, + tiled_api_key: str | None = None, + tiled_prefix: str | None = None, + root_segments: list[str] | None = None, + poll_interval_sec: float = 2.0, + frame_poll_interval_sec: float = 0.5, + unchanged_cycles_threshold: int = 5, + seen_scans_file: str = "/tmp/tiled_listener_seen_scans.txt", + ) -> None: + self.operator = operator + self.tiled_uri = tiled_uri.rstrip("/") + self.tiled_api_key = tiled_api_key + self.tiled_prefix = tiled_prefix + self.root_segments = root_segments or [] + self.poll_interval_sec = poll_interval_sec + self.frame_poll_interval_sec = frame_poll_interval_sec + self.unchanged_cycles_threshold = unchanged_cycles_threshold + self.seen_scans_file = seen_scans_file + self._seen_scans: set[str] = self._load_seen_scans() + + def _load_seen_scans(self) -> set[str]: + """Load previously seen scan UUIDs from disk.""" + try: + with open(self.seen_scans_file) as f: + uuids = {line.strip() for line in f if line.strip()} + logger.info(f"Loaded {len(uuids)} previously seen scan(s) from {self.seen_scans_file}") + return uuids + except FileNotFoundError: + return set() + + def _save_seen_scans(self) -> None: + """Persist seen scan UUIDs to disk.""" + try: + with open(self.seen_scans_file, "w") as f: + f.write("\n".join(self._seen_scans)) + except Exception as e: + logger.warning(f"Could not save seen scans: {e}") + + async def start(self) -> None: + logger.info( + f"TiledNewScanListener started — tiled_uri={self.tiled_uri}, " + f"prefix={self.tiled_prefix}, root_segments={self.root_segments}, " + f"polling every {self.poll_interval_sec}s" + ) + loop = asyncio.get_running_loop() + while True: + try: + await asyncio.to_thread(self._poll_once, loop) + except Exception as e: + logger.exception(f"Error in scan polling loop: {e}") + await asyncio.sleep(self.poll_interval_sec) + + def _get_todays_scan_container(self) -> Container | None: + """Navigate to today's date container, returning None if it doesn't exist yet. + + Reconnects to Tiled on each call to bypass client-side container caching — + without this, newly written UUID containers are invisible to a stale client. + """ + try: + # Fresh client each poll — Tiled caches container children in memory + # so a long-lived client won't see containers created after startup. + container = from_uri(self.tiled_uri, api_key=self.tiled_api_key) + if self.tiled_prefix: + for segment in self.tiled_prefix.split("/"): + if segment: + container = container[segment] + for segment in self.root_segments: + container = container[segment] + now = datetime.now(CALIFORNIA_TZ) + for segment in [str(now.year), f"{now.month:02d}", f"{now.day:02d}"]: + container = container[segment] + return container + except KeyError as e: + logger.debug(f"Scan container path not yet available: {e}") + return None + except Exception as e: + logger.debug(f"Could not resolve scan container: {e}") + return None + + def _poll_once(self, loop: asyncio.AbstractEventLoop) -> None: + """Check for new UUID-level containers under today's date; process each one.""" + scans_container = self._get_todays_scan_container() + if scans_container is None: + return + keys = list(scans_container.keys()) + new_keys = [k for k in keys if k not in self._seen_scans] + if new_keys: + logger.debug(f"Found {len(new_keys)} new scan(s), {len(self._seen_scans)} already seen") + for uuid_key in new_keys: + self._seen_scans.add(uuid_key) + self._save_seen_scans() + try: + scan_container = scans_container[uuid_key] + self._process_scan(uuid_key, scan_container, loop) + except Exception as e: + logger.exception(f"Error processing scan '{uuid_key}': {e}") + + def _process_scan(self, uuid_key: str, scan_container: Container, loop: asyncio.AbstractEventLoop) -> None: + """Stream all frames from a scan container, then emit stop.""" + try: + array_client = scan_container[ARRAY_KEY] + except KeyError: + logger.warning(f"Scan '{uuid_key}' has no '{ARRAY_KEY}' array — skipping") + return + sent_frames: int = 0 + start_sent = False + unchanged_cycles: int = 0 + height: int | None = None + width: int | None = None + + logger.info(f"Processing scan '{uuid_key}'") + + while True: + # Re-fetch array client each cycle — .shape is cached on the object + # so reusing the same client will not see new frames written by Block 1. + try: + array_client = from_uri(self.tiled_uri, api_key=self.tiled_api_key)[ + tuple( + ([s for s in self.tiled_prefix.split("/") if s] if self.tiled_prefix else []) + + self.root_segments + + [str(datetime.now(CALIFORNIA_TZ).year), + f"{datetime.now(CALIFORNIA_TZ).month:02d}", + f"{datetime.now(CALIFORNIA_TZ).day:02d}"] + + [uuid_key, ARRAY_KEY] + ) + ] + except Exception as e: + logger.warning(f"Could not re-fetch array client for '{uuid_key}': {e}") + time.sleep(self.frame_poll_interval_sec) + continue + + n_available = array_client.shape[0] + if height is None: + height, width = array_client.shape[1], array_client.shape[2] + + if not start_sent and n_available > 0: + start_msg = SASStart( + run_name=uuid_key, + run_id=uuid_key, + tiled_url=array_client.uri, + width=width, + height=height, + data_type=str(array_client.dtype), + ) + asyncio.run_coroutine_threadsafe( + self.operator.process(start_msg), loop + ).result() + start_sent = True + + # Emit any new frames + while sent_frames < n_available: + frame_array = array_client[sent_frames] + tiled_url = build_tiled_url( + tiled_uri=self.tiled_uri, + tiled_prefix=self.tiled_prefix, + root_segments=self.root_segments, + uuid=uuid_key, + shot_index=sent_frames, + height=height, + width=width, + ) + raw_event = RawFrameEvent( + image=SerializableNumpyArrayModel(array=frame_array), + frame_number=sent_frames, + tiled_url=tiled_url, + ) + asyncio.run_coroutine_threadsafe( + self.operator.process(raw_event), loop + ).result() + sent_frames += 1 + logger.debug(f"Emitted frame {sent_frames - 1} for scan '{uuid_key}'") + + time.sleep(self.frame_poll_interval_sec) + # Re-fetch after sleep to get the latest frame count from Tiled + try: + refreshed = from_uri(self.tiled_uri, api_key=self.tiled_api_key)[ + tuple( + ([s for s in self.tiled_prefix.split("/") if s] if self.tiled_prefix else []) + + self.root_segments + + [str(datetime.now(CALIFORNIA_TZ).year), + f"{datetime.now(CALIFORNIA_TZ).month:02d}", + f"{datetime.now(CALIFORNIA_TZ).day:02d}"] + + [uuid_key, ARRAY_KEY] + ) + ] + n_after_sleep = refreshed.shape[0] + except Exception: + n_after_sleep = n_available + + if n_after_sleep == sent_frames and n_after_sleep > 0: + unchanged_cycles += 1 + logger.debug(f"No new frames for scan '{uuid_key}' — unchanged cycle {unchanged_cycles}/5") + if unchanged_cycles >= self.unchanged_cycles_threshold: + break + else: + unchanged_cycles = 0 + + if start_sent: + asyncio.run_coroutine_threadsafe( + self.operator.process(SASStop(num_frames=sent_frames)), loop + ).result() + logger.info(f"Scan '{uuid_key}' complete — {sent_frames} frames emitted") + + async def stop(self) -> None: + pass + + async def listen(self) -> None: + pass + + @classmethod + def from_settings(cls, settings, operator: Operator) -> "TiledNewScanListener": + return cls( + operator=operator, + tiled_uri=settings.uri, + tiled_api_key=getattr(settings, "api_key", None), + tiled_prefix=getattr(settings, "tiled_prefix", None), + root_segments=settings.root_segments.to_list(), + poll_interval_sec=getattr(settings, "poll_interval_sec", 2.0), + frame_poll_interval_sec=getattr(settings, "frame_poll_interval_sec", 0.5), + ) + + +def tiled_new_scan_listener_factory( + operator: Operator, + tiled_uri: str, + root_segments: list[str], + tiled_prefix: str | None = None, + poll_interval_sec: float = 2.0, + frame_poll_interval_sec: float = 0.5, + unchanged_cycles_threshold: int = 5, + seen_scans_file: str = "/tmp/tiled_listener_seen_scans.txt", +) -> TiledNewScanListener: + """Factory function for YAML-based wiring.""" + tiled_api_key = os.getenv("RESULTS_TILED_API_KEY") or None + logger.info(f"TiledNewScanListener factory — tiled_uri={tiled_uri}, prefix={tiled_prefix}, root_segments={root_segments}") + return TiledNewScanListener( + operator=operator, + tiled_uri=tiled_uri, + tiled_api_key=tiled_api_key, + tiled_prefix=tiled_prefix, + root_segments=root_segments, + poll_interval_sec=poll_interval_sec, + frame_poll_interval_sec=frame_poll_interval_sec, + unchanged_cycles_threshold=unchanged_cycles_threshold, + seen_scans_file=seen_scans_file, + ) \ No newline at end of file diff --git a/src/tr_ap_xps/xps_tiled_publisher.py b/src/tr_ap_xps/xps_tiled_publisher.py new file mode 100644 index 0000000..45b93dd --- /dev/null +++ b/src/tr_ap_xps/xps_tiled_publisher.py @@ -0,0 +1,328 @@ +import asyncio +import logging +import os +import re +from datetime import datetime +from uuid import uuid4 + +import numpy as np +import pytz +from arroyopy.publisher import Publisher +from tiled.client import from_uri +from tiled.client.array import ArrayClient +from tiled.client.node import Container + +from .schemas import XPSResult, XPSResultStart, XPSResultStop + +logger = logging.getLogger(__name__) + +LOCAL_TILED_API_KEY = os.getenv("RESULTS_TILED_API_KEY") or None +CALIFORNIA_TZ = pytz.timezone("US/Pacific") +UUID_PATTERN = re.compile( + r"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}" +) +ARRAY_KEY = "xps_averaged_heatmaps" + + +def extract_uuid(scan_name: str) -> str | None: + """Extract a UUID36 from a scan_name string. + + Args: + scan_name: Scan name that may contain a UUID36 substring. + + Returns: + The matched UUID string, or None if not found. + """ + match = UUID_PATTERN.search(scan_name) + return match.group() if match else None + + +class XPSTiledResultPublisher(Publisher): + """Publisher that writes XPS shot_mean frames to local Tiled. + + Container path structure: + [tiled_prefix /] / / /
/ / xps_averaged_heatmaps + + Args: + tiled_uri: Base URI of the Tiled server (e.g. ``"http://tiled:8000"``). + tiled_api_key: API key for Tiled authentication. Falls back to the + ``RESULTS_TILED_API_KEY`` environment variable when omitted. + tiled_prefix: Top-level container key under which all data is written. + root_segments: Intermediate path segments between prefix and date + (e.g. ``["xps_processed_images"]``). + """ + + def __init__( + self, + tiled_uri: str, + tiled_api_key: str | None = None, + tiled_prefix: str | None = None, + root_segments: list[str] | None = None, + ) -> None: + super().__init__() + self.tiled_uri = tiled_uri + self.tiled_api_key = tiled_api_key or LOCAL_TILED_API_KEY + self.tiled_prefix = tiled_prefix + self.root_segments = root_segments or [] + + # Tiled state + self._tiled_client: Container | None = None + self._array_clients: dict[str, ArrayClient] = {} + + # Per-scan state + self._current_uuid: str | None = None + self._current_scan_name: str | None = None + self._shot_index: int = 0 + self._sas_start_sent: bool = False + self._scan_started: bool = False # True only after XPSResultStart received + + logger.info( + f"Initialized XPSTiledResultPublisher — tiled_uri={tiled_uri}, " + f"prefix={tiled_prefix}, root_segments={self.root_segments}" + ) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + async def start(self) -> None: + """Connect to the Tiled server eagerly so errors surface at startup.""" + await asyncio.to_thread(self._connect_tiled) + + def _connect_tiled(self) -> None: + """Open a synchronous Tiled client connection.""" + logger.info(f"Connecting to Tiled server at {self.tiled_uri}") + self._tiled_client = from_uri(self.tiled_uri, api_key=self.tiled_api_key) + + # ------------------------------------------------------------------ + # Publisher entry point + # ------------------------------------------------------------------ + + async def publish( + self, message: XPSResultStart | XPSResult | XPSResultStop + ) -> None: + """Route incoming messages to the appropriate handler. + + Args: + message: One of XPSResultStart, XPSResult, or XPSResultStop. + """ + if isinstance(message, XPSResultStart): + await self._handle_start(message) + elif isinstance(message, XPSResult): + await self._handle_event(message) + elif isinstance(message, XPSResultStop): + await self._handle_stop(message) + else: + logger.warning(f"Unhandled message type: {type(message)}") + + # ------------------------------------------------------------------ + # Message handlers + # ------------------------------------------------------------------ + + async def _handle_start(self, message: XPSResultStart) -> None: + """Reset per-scan state. + + Args: + message: The XPSResultStart message. + """ + self._current_scan_name = message.scan_name + self._current_uuid = extract_uuid(message.scan_name or "") + self._array_clients = {} + self._shot_index = 0 + self._sas_start_sent = False + self._scan_started = True # mark scan as active; frames before this are discarded + + if self._current_uuid is None: + self._current_uuid = str(uuid4()) + logger.warning( + f"No UUID found in scan_name='{message.scan_name}' — " + f"generated new UUID: '{self._current_uuid}'" + ) + else: + logger.info( + f"Start received — scan_name='{self._current_scan_name}', " + f"uuid='{self._current_uuid}'" + ) + + async def _handle_event(self, message: XPSResult) -> None: + """Write shot_mean to Tiled. + + Args: + message: The XPSResult message containing shot_mean. + """ + if not self._scan_started: + logger.warning( + f"Received scan '{self._current_scan_name}' frame {message.frame_number} " + "but no start message has been received yet — discarding frame to avoid " + "Tiled index desync." + ) + return + + if message.shot_mean is None: + logger.debug("shot_mean is None — skipping frame") + return + + if self._current_uuid is None: + logger.error("No valid UUID for current scan — ignoring frame") + return + + shot_array: np.ndarray = message.shot_mean.array + + if shot_array.ndim != 2: + logger.warning( + f"Expected 2-D shot_mean, got shape {shot_array.shape} — skipping" + ) + return + + # --- Write to Tiled --- + written_index = self._shot_index # capture before incrementing + try: + array_client = await asyncio.to_thread( + self._get_or_create_array_client, self._current_uuid, shot_array + ) + if array_client is None: + logger.warning(f"Failed to get array client for frame {message.frame_number}") + return + + # Frame 0 was already written by _get_or_create_array_client via write_array + if written_index > 0: + await asyncio.to_thread(self._patch_array, array_client, shot_array) + logger.debug( + f"Patched {ARRAY_KEY} for uuid='{self._current_uuid}', " + f"shot_index={written_index}" + ) + else: + logger.info(f"Initialised {ARRAY_KEY} for uuid='{self._current_uuid}'") + + self._shot_index += 1 + except Exception as e: + logger.error( + f"Error writing shot_mean for frame {message.frame_number}: {e}", + exc_info=True, + ) + + async def _handle_stop(self, message: XPSResultStop) -> None: + """Clear per-scan state. + + Args: + message: The XPSResultStop message. + """ + logger.info( + f"Stop received — clearing state for scan='{self._current_scan_name}', " + f"uuid='{self._current_uuid}'" + ) + + self._array_clients = {} + self._current_scan_name = None + self._current_uuid = None + self._shot_index = 0 + self._sas_start_sent = False + self._scan_started = False # reset; next frames must wait for a new start + + # ------------------------------------------------------------------ + # Tiled helpers (synchronous — called via asyncio.to_thread) + # ------------------------------------------------------------------ + + def _get_or_create_array_client( + self, uuid: str, first_shot: np.ndarray + ) -> ArrayClient | None: + """Get a cached array client for uuid, or create the container hierarchy and + write the initial frame if this is the first time we've seen this uuid. + + Path: [tiled_prefix /] / / / xps_averaged_heatmaps + + Args: + uuid: UUID extracted from scan_name. + first_shot: 2-D shot_mean array of shape ``(height, width)``. Only used + on the first call for this uuid to initialise the array shape. + + Returns: + ArrayClient pointing at the array, or None on error. + """ + if uuid in self._array_clients: + return self._array_clients[uuid] + + if self._tiled_client is None: + logger.error("Tiled client is not connected. Call start() first.") + return None + + try: + container: Container = self._tiled_client + + if self.tiled_prefix: + for segment in self.tiled_prefix.split("/"): + if segment: + container = self._get_or_create_container(segment, container) + + for segment in self.root_segments: + container = self._get_or_create_container(segment, container) + + now = datetime.now(CALIFORNIA_TZ) + for segment in [str(now.year), f"{now.month:02d}", f"{now.day:02d}"]: + container = self._get_or_create_container(segment, container) + + container = self._get_or_create_container(uuid, container) + + initial_array = first_shot[None, :, :] + logger.info( + f"Creating {ARRAY_KEY} for uuid='{uuid}' with shape {initial_array.shape}" + ) + array_client = container.write_array(initial_array, key=ARRAY_KEY) + self._array_clients[uuid] = array_client + return array_client + + except Exception as e: + logger.error( + f"Error creating array client for uuid='{uuid}': {e}", exc_info=True + ) + return None + + def _get_or_create_container(self, key: str, parent: Container) -> Container: + """Return the child container at key, creating it if absent. + + Args: + key: Container key to look up or create. + parent: Parent Tiled container. + + Returns: + The existing or newly-created child container. + """ + if key in parent: + return parent[key] + logger.info(f"Creating container: {key}") + return parent.create_container(key) + + def _patch_array(self, array_client: ArrayClient, array: np.ndarray) -> None: + """Append a new 2-D frame to the existing 3-D Tiled array. + + Args: + array_client: The ArrayClient to patch into. + array: 2-D shot_mean array of shape ``(height, width)``. + """ + current_shape = array_client.shape # (n_shots, height, width) + offset = (current_shape[0],) + logger.debug(f"Patching at offset {offset}, current shape {current_shape}") + array_client.patch(array[None, :, :], offset=offset, extend=True) + logger.debug(f"Patch complete — new shape {array_client.shape}") + + +def xps_tiled_result_publisher_factory( + tiled_uri: str, + tiled_prefix: str | None = None, + root_segments: list[str] | None = None, +) -> XPSTiledResultPublisher: + """Instantiate XPSTiledResultPublisher for YAML-based wiring. + + Args: + tiled_uri: Base URI of the Tiled server. + tiled_prefix: Optional top-level container prefix. + root_segments: Intermediate path segments between prefix and date. + + Returns: + A configured XPSTiledResultPublisher instance. + """ + return XPSTiledResultPublisher( + tiled_uri=tiled_uri, + tiled_prefix=tiled_prefix, + root_segments=root_segments, + ) \ No newline at end of file diff --git a/src/tr_ap_xps/zmq_publisher.py b/src/tr_ap_xps/zmq_publisher.py new file mode 100644 index 0000000..99fe92f --- /dev/null +++ b/src/tr_ap_xps/zmq_publisher.py @@ -0,0 +1,45 @@ +import logging + +import msgpack +import zmq +import zmq.asyncio +from arroyopy.publisher import Publisher +from zmq.asyncio import Context, Socket + +from arroyosas.schemas import RawFrameEvent, SASMessage, SASStart, SASStop + +logger = logging.getLogger(__name__) + + +class ZMQFramePublisher(Publisher): + def __init__(self, zmq_socket: Socket): + self.zmq_socket = zmq_socket + + async def publish(self, message: SASMessage) -> None: + logger.debug(f"Publishing message: {message.msg_type}") + if isinstance(message, SASStart) or isinstance(message, SASStop): + message = msgpack.packb(message.model_dump(), use_bin_type=True) + await self.zmq_socket.send(message) + return + if isinstance(message, RawFrameEvent): + message = message.model_dump() + message = msgpack.packb(message, use_bin_type=True) + await self.zmq_socket.send(message) + else: + logger.warning(f"Unknown message type: {type(message)}") + + @classmethod + def from_settings(cls, settings) -> "ZMQFramePublisher": + context = Context() + zmq_socket = context.socket(zmq.PUB) + zmq_socket.bind(settings.address) + logger.info(f"##### Publishing frames to {settings.address}") + return cls(zmq_socket) + + +def create_zmq_frame_publisher(zmq_address: str) -> ZMQFramePublisher: + context = Context() + zmq_socket = context.socket(zmq.PUB) + zmq_socket.bind(zmq_address) + logger.info(f"##### Publishing frames to {zmq_address}") + return ZMQFramePublisher(zmq_socket)