fastcs-pandablocks 0.2.0a3__py3-none-any.whl

fastcs_pandablocks/panda/blocks/data.py

@@ -0,0 +1,625 @@
+import asyncio
+import enum
+import logging
+import os
+from asyncio import CancelledError
+from collections import deque
+from collections.abc import AsyncGenerator, Callable
+from dataclasses import dataclass
+from importlib.util import find_spec
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+from fastcs.attributes import AttrR, AttrRW
+from fastcs.controllers import Controller
+from fastcs.datatypes import Bool, Enum, Float, Int, String, Table
+from numpy.typing import DTypeLike
+from pandablocks.hdf import (
+    EndData,
+    FrameData,
+    Pipeline,
+    StartData,
+    create_default_pipeline,
+    stop_pipeline,
+)
+from pandablocks.responses import Data, EndReason, ReadyData
+
+from fastcs_pandablocks.types import PandaName
+
+HDFReceived = Union[ReadyData, StartData, FrameData, EndData]
+
+
+class CaptureMode(enum.Enum):
+    """
+    The mode which the circular buffer will use to flush.
+    """
+
+    #: Wait till N frames are recieved then write them
+    #:  and finish capture
+    FIRST_N = 0
+
+    #: On EndData write the last N frames
+    LAST_N = 1
+
+    #: Write data as received until Capture set to 0
+    FOREVER = 2
+
+
+class NumCapturedSetter(Pipeline):
+    def __init__(self, number_captured_setter: Callable) -> None:
+        self.number_captured_setter = number_captured_setter
+        super().__init__()
+
+        # TODO: Sync expected but async passed. Fix.
+        self.what_to_do = {int: self.set_record}
+
+    def set_record(self, value: int):
+        asyncio.run(self.number_captured_setter(value))
+
+
+class HDF5Buffer:
+    _buffer_index = None
+    start_data = None
+    number_of_received_rows = 0
+    finish_capturing = False
+    number_of_rows_in_circular_buffer = 0
+
+    def __init__(
+        self,
+        capture_mode: CaptureMode,
+        filepath: Path,
+        number_of_rows_to_capture: int,
+        status_message_setter: Callable,
+        number_received_setter: Callable,
+        number_captured_setter_pipeline: NumCapturedSetter,
+        dataset_name_cache: dict[str, dict[str, str]],
+    ):
+        # Only one filename - user must stop capture and set new FileName/FilePath
+        # for new files
+
+        self.circular_buffer: deque[FrameData] = deque()
+        self.capture_mode = capture_mode
+
+        match capture_mode:
+            case CaptureMode.FIRST_N:
+                self._handle_FrameData = self._capture_first_n
+            case CaptureMode.LAST_N:
+                self._handle_FrameData = self._capture_last_n
+            case CaptureMode.FOREVER:
+                self._handle_FrameData = self._capture_forever
+            case _:
+                raise RuntimeError("Invalid capture mode")
+
+        self.filepath = filepath
+        self.number_of_rows_to_capture = number_of_rows_to_capture
+        self.status_message_setter = status_message_setter
+        self.number_received_setter = number_received_setter
+        self.number_captured_setter_pipeline = number_captured_setter_pipeline
+
+        self.dataset_name_cache = dataset_name_cache
+
+        if (
+            self.capture_mode == CaptureMode.LAST_N
+            and self.number_of_rows_to_capture <= 0
+        ):
+            raise RuntimeError("Number of rows to capture must be > 0 on LAST_N mode")
+
+        self.start_pipeline()
+
+    def __del__(self):
+        if self.pipeline[0].is_alive():
+            stop_pipeline(self.pipeline)
+
+    def put_data_to_file(self, data: HDFReceived):
+        try:
+            self.pipeline[0].queue.put_nowait(data)
+        except Exception as ex:
+            logging.exception(f"Failed to save the data to HDF5 file: {ex}")
+
+    def start_pipeline(self):
+        self.pipeline = create_default_pipeline(
+            iter([str(self.filepath)]),
+            self.dataset_name_cache,
+            self.number_captured_setter_pipeline,
+        )
+
+    def _handle_StartData(self, data: StartData):
+        if self.start_data and data != self.start_data:
+            # PandA was disarmed, had config changed, and rearmed.
+            # Cannot process to the same file with different start data.
+            logging.error(
+                "New start data detected, differs from previous start "
+                "data for this file. Aborting HDF5 data capture."
+            )
+
+            self.status_message_setter(
+                "Mismatched StartData packet for file",
+            )
+            self.put_data_to_file(
+                EndData(self.number_of_received_rows, EndReason.START_DATA_MISMATCH)
+            )
+
+            self.finish_capturing = True
+
+        # Only pass StartData to pipeline if we haven't previously
+        else:
+            # In LAST_N mode, wait till the end of capture to write
+            # the StartData to file.
+            # In FOREVER mode write the StartData to file if it's the first received.
+            if (
+                self.capture_mode == CaptureMode.FIRST_N
+                or self.capture_mode == CaptureMode.FOREVER
+                and not self.start_data
+            ):
+                self.put_data_to_file(data)
+
+            self.start_data = data
+
+    async def _capture_first_n(self, data: FrameData):
+        """
+        Capture framedata as it comes in. Stop when number of frames exceeds
+        number_of_rows_to_capture, and cut off the data so that it's length
+        number_of_rows_to_capture.
+        """
+        self.number_of_received_rows += len(data.data)
+
+        if (
+            self.number_of_rows_to_capture > 0
+            and self.number_of_received_rows > self.number_of_rows_to_capture
+        ):
+            # Discard extra collected data points if necessary
+            data.data = data.data[
+                : self.number_of_rows_to_capture - self.number_of_received_rows
+            ].copy()
+            self.number_of_received_rows = self.number_of_rows_to_capture
+
+        self.put_data_to_file(data)
+        await self.number_received_setter(self.number_of_received_rows)
+
+        if (
+            self.number_of_rows_to_capture > 0
+            and self.number_of_received_rows == self.number_of_rows_to_capture
+        ):
+            # Reached configured capture limit, stop the file
+            logging.info(
+                f"Requested number of frames ({self.number_of_rows_to_capture}) "
+                "captured, disabling Capture."
+            )
+            self.status_message_setter("Requested number of frames captured")
+            self.put_data_to_file(EndData(self.number_of_received_rows, EndReason.OK))
+            self.finish_capturing = True
+
+    async def _capture_forever(self, data: FrameData):
+        self.put_data_to_file(data)
+        self.number_of_received_rows += len(data.data)
+        await self.number_received_setter(self.number_of_received_rows)
+
+    async def _capture_last_n(self, data: FrameData):
+        """
+        Append every FrameData to a buffer until the number of rows equals
+        `:NumCapture`. Then rewrite the data circularly.
+
+        Only write the data once PCAP is received.
+        """
+        self.circular_buffer.append(data)
+        self.number_of_received_rows += len(data.data)
+        self.number_of_rows_in_circular_buffer += len(data.data)
+
+        if self.number_of_rows_in_circular_buffer > self.number_of_rows_to_capture:
+            await self.status_message_setter(
+                "NumCapture received, rewriting first frames received"
+            )
+
+        else:
+            await self.status_message_setter("Filling buffer to NumReceived")
+
+        while self.number_of_rows_in_circular_buffer > self.number_of_rows_to_capture:
+            first_frame_data = self.circular_buffer.popleft()
+            first_frame_data_length = len(first_frame_data.data)
+
+            if first_frame_data_length > self.number_of_rows_to_capture:
+                # More data than we want to capture, all in a single FrameData
+                # We can just slice with the NumCapture since this has to be the
+                # only FrameData in the buffer at this point
+                assert len(self.circular_buffer) == 0
+                shrinked_data = first_frame_data.data[
+                    -self.number_of_rows_to_capture :
+                ].copy()
+                first_frame_data.data = shrinked_data
+                self.circular_buffer.appendleft(first_frame_data)
+                self.number_of_rows_in_circular_buffer = self.number_of_rows_to_capture
+            elif (
+                first_frame_data_length
+                > self.number_of_rows_in_circular_buffer
+                - self.number_of_rows_to_capture
+            ):
+                # We can slice from the beginning of the FrameData to have the desired
+                # number of rows
+                indices_to_discard = (
+                    self.number_of_rows_in_circular_buffer
+                    - self.number_of_rows_to_capture
+                )
+                shrinked_data = first_frame_data.data[indices_to_discard:].copy()
+                first_frame_data.data = shrinked_data
+                self.circular_buffer.appendleft(first_frame_data)
+                self.number_of_rows_in_circular_buffer -= indices_to_discard
+                assert (
+                    self.number_of_rows_in_circular_buffer
+                    == self.number_of_rows_to_capture
+                )
+            else:
+                # If we remove the enire first frame data then the buffer will still
+                # be too big, or it will be exactly the number of rows we want
+                self.number_of_rows_in_circular_buffer -= first_frame_data_length
+
+        await self.number_received_setter(self.number_of_received_rows)
+
+    def _handle_EndData(self, data: EndData):
+        match self.capture_mode:
+            case CaptureMode.LAST_N:
+                # In LAST_N only write FrameData if the EndReason is OK
+                if data.reason not in (EndReason.OK, EndReason.MANUALLY_STOPPED):
+                    self.status_message_setter(
+                        f"Stopped capturing with reason {data.reason}, "
+                        "skipping writing of buffered frames"
+                    )
+                    self.finish_capturing = True
+                    return
+
+                self.status_message_setter(
+                    "Finishing capture, writing buffered frames to file"
+                )
+                assert self.start_data is not None
+                self.put_data_to_file(self.start_data)
+                for frame_data in self.circular_buffer:
+                    self.put_data_to_file(frame_data)
+
+            case CaptureMode.FOREVER:
+                if data.reason != EndReason.MANUALLY_STOPPED:
+                    self.status_message_setter(
+                        "Finished capture, waiting for next ReadyData"
+                    )
+                    return
+
+            case CaptureMode.FIRST_N:
+                pass  # Frames will have already been written in FirstN
+
+            case _:
+                raise RuntimeError("Unknown capture mode")
+
+        self.status_message_setter("Finished capture")
+        self.finish_capturing = True
+        self.put_data_to_file(data)
+
+    async def handle_data(self, data: HDFReceived):
+        match data:
+            case ReadyData():
+                pass
+            case StartData():
+                await self.status_message_setter("Starting capture")
+                self._handle_StartData(data)
+            case FrameData():
+                await self._handle_FrameData(data)
+            case EndData():
+                self._handle_EndData(data)
+            case _:
+                raise RuntimeError(
+                    f"Data was recieved that was of type {type(data)}, not"
+                    "StartData, EndData, ReadyData, or FrameData"
+                )
+
+
+@dataclass
+class DatasetAttributes:
+    """A dataset name and capture mode"""
+
+    name: AttrRW[str]
+    capture: AttrRW[enum.Enum]
+
+
+class DatasetTableWrapper:
+    """Used for outputing formatted dataset names in the HDF5 writer, and creating
+    and updating the HDF5 `DATASETS` table attribute."""
+
+    NUMPY_TYPE: list[tuple[str, DTypeLike]] = [
+        ("name", np.dtype("S1000")),
+        ("dtype", np.dtype("S1000")),
+    ]
+
+    def __init__(
+        self,
+        dataset_cache: dict[PandaName, DatasetAttributes],
+    ):
+        self._dataset_cache = dataset_cache
+
+    def hdf_writer_names(self) -> dict[str, dict[str, str]]:
+        """Formats the current dataset names for use in the HDFWriter"""
+
+        hdf_names: dict[str, dict[str, str]] = {}
+        for panda_name, dataset in self._dataset_cache.items():
+            capture_str_value = dataset.capture.get().value
+            name_str_value = dataset.name.get()
+            if not name_str_value or capture_str_value == "No":
+                continue
+
+            hdf_names[str(panda_name)] = hdf_name = {}
+
+            hdf_name[capture_str_value.split(" ")[-1]] = name_str_value
+            # Suffix -min and -max if both are present
+            if "Min Max" in capture_str_value:
+                hdf_name["Min"] = f"{name_str_value}-min"
+                hdf_name["Max"] = f"{name_str_value}-max"
+
+        return hdf_names
+
+    def get_numpy_table(self) -> np.ndarray:
+        return np.array(
+            [
+                (dataset.name.get(), "float64")
+                for dataset in self._dataset_cache.values()
+                if dataset.name.get() and dataset.capture.get() != "No"
+            ],
+            dtype=self.NUMPY_TYPE,
+        )
+
+    def set_on_update_callback(self, table_attribute: AttrR):
+        async def callback(value):
+            await table_attribute.update(self.get_numpy_table())
+
+        for dataset_attributes in self._dataset_cache.values():
+            dataset_attributes.name.set_update_callback(callback)
+            dataset_attributes.capture.set_update_callback(callback)
+
+
+class DataController(Controller):
+    """Class to create and control the records that handle HDF5 processing"""
+
+    hdf_directory = AttrRW(String(), description="File path for HDF5 files.")
+
+    create_directory = AttrRW(
+        Int(),
+        description="Directory creation depth",
+        initial_value=0,
+    )
+
+    directory_exists = AttrR(
+        Bool(), description="Directory exists", initial_value=False
+    )
+
+    hdf_file_name = AttrRW(
+        String(),
+        description="File name prefix for HDF5 files",
+        initial_value="",
+    )
+
+    hdf_full_file_path = AttrR(
+        String(),
+        description="Full HDF5 file name with directory",
+        initial_value="",
+    )
+
+    num_capture = AttrRW(
+        Int(min=0),
+        description="Number of frames to capture. 0=infinite",
+        initial_value=0,  # Infinite capture
+    )
+
+    num_captured = AttrR(
+        Int(),
+        description="Number of frames written to file.",
+        initial_value=0,
+    )
+
+    num_received = AttrR(
+        Int(),
+        description="Number of frames received from panda.",
+        initial_value=0,
+    )
+    flush_period = AttrRW(
+        Float(units="s"),
+        description="Frequency that data is flushed (seconds).",
+        initial_value=1.0,
+    )
+
+    capture = AttrRW(
+        Bool(), description="Start/stop HDF5 capture.", initial_value=False
+    )
+    capture_mode = AttrRW(
+        Enum(CaptureMode),
+        description="Choose how to hdf writer flushes",
+        initial_value=CaptureMode.FIRST_N,
+    )
+
+    status = AttrR(
+        String(),
+        description="Status of HDF5 capture",
+        initial_value="OK",
+    )
+
+    def __init__(
+        self,
+        client_data: Callable[[bool, float], AsyncGenerator[Data, None]],
+        dataset_attributes: dict[PandaName, DatasetAttributes],
+    ):
+        super().__init__()
+
+        if find_spec("h5py") is None:
+            logging.warning("No HDF5 support detected - skipping creating HDF5 records")
+            return
+
+        self._client_data = client_data
+        self._dataset_table_wrapper = DatasetTableWrapper(dataset_attributes)
+        self._handle_hdf5_data_task = None
+
+        datasets_attribute = AttrR(
+            Table(self._dataset_table_wrapper.NUMPY_TYPE),
+            description="HDF5 dataset names.",
+            initial_value=self._dataset_table_wrapper.get_numpy_table(),
+        )
+        self.attributes["datasets"] = datasets_attribute
+        self._dataset_table_wrapper.set_on_update_callback(datasets_attribute)
+
+        self.hdf_directory.add_on_update_callback(self._update_directory_path)
+        self.hdf_file_name.add_on_update_callback(self._update_full_file_path)
+        self.capture.add_on_update_callback(self._capture_on_update)
+
+    async def _update_directory_path(self, new_val) -> None:
+        """Handles writes to the directory path PV, creating
+        directories based on the setting of the CreateDirectory record"""
+        new_path = Path(new_val).absolute()
+        create_dir_depth = self.create_directory.get()
+        max_dirs_to_create = 0
+        if create_dir_depth < 0:
+            max_dirs_to_create = abs(create_dir_depth)
+        elif create_dir_depth > len(new_path.parents):
+            max_dirs_to_create = 0
+        elif create_dir_depth > 0:
+            max_dirs_to_create = len(new_path.parents) - create_dir_depth
+
+        logging.debug(f"Permitted to create up to {max_dirs_to_create} dirs.")
+        dirs_to_create = 0
+        for p in reversed(new_path.parents):
+            if not p.exists():
+                if dirs_to_create == 0:
+                    # First directory level that does not exist, log it.
+                    logging.error(f"All dir from {str(p)} and below do not exist!")
+                dirs_to_create += 1
+            else:
+                logging.info(f"{str(p)} exists")
+
+        # Account for target path itself not existing
+        if not os.path.exists(new_path):
+            dirs_to_create += 1
+
+        logging.debug(f"Need to create {dirs_to_create} directories.")
+
+        # Case where all dirs exist
+        if dirs_to_create == 0:
+            if os.access(new_path, os.W_OK):
+                status_msg = "Dir exists and is writable"
+                await self.directory_exists.update(True)
+            else:
+                status_msg = "Dirs exist but aren't writable."
+                await self.directory_exists.update(False)
+        # Case where we will create directories
+        elif dirs_to_create <= max_dirs_to_create:
+            logging.debug(f"Attempting to create {dirs_to_create} dir(s)...")
+            try:
+                os.makedirs(new_path, exist_ok=True)
+                status_msg = f"Created {dirs_to_create} dirs."
+                await self.directory_exists.update(True)
+            except PermissionError:
+                status_msg = "Permission error creating dirs!"
+                await self.directory_exists.update(False)
+        # Case where too many directories need to be created
+        else:
+            status_msg = f"Need to create {dirs_to_create} > {max_dirs_to_create} dirs."
+            await self.directory_exists.update(False)
+
+        if self.directory_exists.get() == 0:
+            logging.error(status_msg)
+        else:
+            logging.debug(status_msg)
+
+        await self.status.update(status_msg)
+
+        await self._update_full_file_path(new_val)
+
+    async def _update_full_file_path(self, value) -> None:
+        await self.hdf_full_file_path.update(self._get_filepath())
+
+    async def _handle_hdf5_data(self) -> None:
+        """Handles writing HDF5 data from the PandA to file, based on configuration
+        in the various HDF5 records.
+        This method expects to be run as an asyncio Task."""
+        buffer: HDF5Buffer | None = None
+        try:
+            # Set up the hdf buffer
+
+            # TODO: Check if exists or writeable
+            if self.hdf_directory.get() == "":
+                raise RuntimeError(
+                    "Configured HDF directory does not exist or is not writable!"
+                )
+
+            num_capture: int = self.num_capture.get()
+            capture_mode: CaptureMode = CaptureMode(self.capture_mode.get())
+            filepath = self._get_filepath()
+
+            await self.num_captured.update(0)
+            number_captured_setter_pipeline = NumCapturedSetter(
+                self.num_captured.update
+            )
+
+            await self.attributes["datasets"].update(  # type: ignore
+                self._dataset_table_wrapper.get_numpy_table()
+            )
+
+            buffer = HDF5Buffer(
+                capture_mode,
+                Path(filepath),
+                num_capture,
+                self.status.update,
+                self.num_received.update,
+                number_captured_setter_pipeline,
+                self._dataset_table_wrapper.hdf_writer_names(),
+            )
+            flush_period: float = self.flush_period.get()
+            async for data in self._client_data(False, flush_period):
+                logging.debug(f"Received data packet: {data}")
+
+                await buffer.handle_data(data)  # type: ignore
+                if buffer.finish_capturing:
+                    break
+
+        except CancelledError:
+            logging.info("Capturing task cancelled, closing HDF5 file")
+            await self.status.update("Capturing disabled")
+            # Only send EndData if we know the file was opened - could be cancelled
+            # before PandA has actually send any data
+            if buffer and buffer.capture_mode != CaptureMode.LAST_N:
+                buffer.put_data_to_file(
+                    EndData(buffer.number_of_received_rows, EndReason.MANUALLY_STOPPED)
+                )
+
+        except Exception:
+            logging.exception("HDF5 data capture terminated due to unexpected error")
+            await self.status.update(
+                "Capture disabled, unexpected exception.",
+            )
+            # Only send EndData if we know the file was opened - exception could happen
+            # before file was opened
+            if (
+                buffer
+                and buffer.start_data
+                and buffer.capture_mode != CaptureMode.LAST_N
+            ):
+                buffer.put_data_to_file(
+                    EndData(buffer.number_of_received_rows, EndReason.UNKNOWN_EXCEPTION)
+                )
+
+        finally:
+            logging.debug("Finishing processing HDF5 PandA data")
+            await self.num_received.update(
+                buffer.number_of_received_rows if buffer else 0
+            )
+            await self.capture.update(False)
+
+    def _get_filepath(self) -> str:
+        """Create the file path for the HDF5 file from the relevant records"""
+        return "/".join([self.hdf_directory.get(), self.hdf_file_name.get()])
+
+    async def _capture_on_update(self, value) -> None:
+        """Process an update to the Capture record, to start/stop recording HDF5 data"""
+        logging.debug(f"Entering HDF5:Capture record on_update method, value {value}.")
+        if value:
+            if self._handle_hdf5_data_task:
+                logging.warning("Existing HDF5 capture running, cancelling it.")
+                self._handle_hdf5_data_task.cancel()
+
+            self._handle_hdf5_data_task = asyncio.create_task(self._handle_hdf5_data())
+        else:
+            if self._handle_hdf5_data_task is not None:
+                self._handle_hdf5_data_task.cancel()  # Abort any HDF5 file writing
+                self._handle_hdf5_data_task = None

fastcs_pandablocks/panda/blocks/versions.py

@@ -0,0 +1,69 @@
+import logging
+import re
+from typing import TypedDict
+
+from fastcs.attributes import AttrR
+from fastcs.controllers import Controller
+from fastcs.datatypes import String
+
+from fastcs_pandablocks.types import WidgetGroup
+
+
+class PandaVersions(TypedDict):
+    panda_sw: str
+    fpga: str
+    rootfs: str
+
+
+def _parse_idn_response(idn_response: str) -> PandaVersions:
+    """Function that parses version info from the PandA's response to the IDN command
+
+    See: https://pandablocks-server.readthedocs.io/en/latest/commands.html#system-commands
+    """
+
+    # Currently, IDN reports sw, fpga, and rootfs versions
+    firmware_versions = {"PandA SW": "Unknown", "FPGA": "Unknown", "rootfs": "Unknown"}
+
+    # If the *IDN response contains too many keys, break and leave versions as "Unknown"
+    # Since spaces are used to deliminate versions and can also be in the keys and
+    # values, if an additional key is present that we don't explicitly handle,
+    # our approach of using regex matching will not work.
+    if sum(name in idn_response for name in firmware_versions) < idn_response.count(
+        ":"
+    ):
+        logging.error(
+            f"Recieved unexpected version numbers in version string {idn_response}!"
+        )
+    else:
+        for firmware_name in firmware_versions:
+            pattern = re.compile(
+                rf"{re.escape(firmware_name)}:\s*([^:]+?)(?=\s*\b(?: \
+                {'|'.join(map(re.escape, firmware_versions))}):|$)"
+            )
+            if match := pattern.search(idn_response):
+                firmware_versions[firmware_name] = match.group(1).strip()
+                logging.info(
+                    f"{firmware_name} Version: {firmware_versions[firmware_name]}"
+                )
+            else:
+                logging.warning(f"Failed to get {firmware_name} version information!")
+
+    return PandaVersions(
+        panda_sw=firmware_versions["PandA SW"],
+        fpga=firmware_versions["FPGA"],
+        rootfs=firmware_versions["rootfs"],
+    )
+
+
+class VersionController(Controller):
+    def __init__(self, idn_response: str):
+        super().__init__()
+        self.description = "Version information from the PandA."
+        versions = _parse_idn_response(idn_response)
+        for version_name, version in versions.items():
+            self.attributes[version_name] = AttrR(
+                String(),
+                description="Version information from the PandA.",
+                group=WidgetGroup.READBACKS.value,
+                initial_value=version,  # type: ignore
+            )