quelware-client 0.1.0__py3-none-any.whl

quelware_client/__init__.py

@@ -0,0 +1,3 @@
+from .client import create_quelware_client
+
+__all__ = ["create_quelware_client"]

quelware_client/client/__init__.py

@@ -0,0 +1,4 @@
+from ._grpc import create_quelware_client
+from ._standalone_grpc import create_standalone_client
+
+__all__ = ["create_quelware_client", "create_standalone_client"]

quelware_client/client/_grpc.py

@@ -0,0 +1,61 @@
+from grpclib.client import Channel
+from quelware_core.entities.unit import UnitLabel
+
+from quelware_client.core import AgentContainer, AgentFactory, QuelwareClient
+from quelware_client.core.interfaces.instrument_agent import InstrumentAgent
+from quelware_client.core.interfaces.resource_agent import ResourceAgent
+from quelware_client.core.interfaces.session_agent import SessionAgent
+from quelware_client.core.interfaces.system_configuration_agent import (
+    SystemConfigurationAgent,
+)
+from quelware_client.infra.instrument_agent_grpc import InstrumentAgentGrpc
+from quelware_client.infra.resource_agent_grpc import (
+    ResourceAgentGrpc,
+)
+from quelware_client.infra.session_agent_grpc import SessionAgentGrpc
+from quelware_client.infra.system_configuration_agent_grpc import (
+    SystemConfigurationAgentGrpc,
+)
+
+
+def _create_default_resource_agent_factory(channel):
+    def _default_resource_agent_factory(ul: UnitLabel):
+        return ResourceAgentGrpc(channel, metadata={"x-unit-label": str(ul)})
+
+    return _default_resource_agent_factory
+
+
+def _create_default_instrument_agent_factory(channel):
+    def _default_command_agent_factory(ul: UnitLabel):
+        return InstrumentAgentGrpc(channel, metadata={"x-unit-label": str(ul)})
+
+    return _default_command_agent_factory
+
+
+def create_quelware_client(  # noqa: PLR0913
+    endpoint: str = "localhost",
+    port: int = 50051,
+    session_agent: SessionAgent | None = None,
+    configuration_agent: SystemConfigurationAgent | None = None,
+    resource_agent_factory: AgentFactory[ResourceAgent] | None = None,
+    instrument_agent_factory: AgentFactory[InstrumentAgent] | None = None,
+):
+    channel = Channel(endpoint, port)
+
+    if resource_agent_factory is None:
+        resource_agent_factory = _create_default_resource_agent_factory(channel)
+
+    if instrument_agent_factory is None:
+        instrument_agent_factory = _create_default_instrument_agent_factory(channel)
+
+    agent_container = AgentContainer()
+    if session_agent is None:
+        agent_container.session = SessionAgentGrpc(channel)
+    if configuration_agent is None:
+        agent_container.system_configuration = SystemConfigurationAgentGrpc(channel)
+    return QuelwareClient(
+        agent=agent_container,
+        resource_agent_factory=resource_agent_factory,
+        instrument_agent_factory=instrument_agent_factory,
+        close_handlers=[channel.close],
+    )

quelware_client/client/_standalone_grpc.py

@@ -0,0 +1,48 @@
+import logging
+
+from grpclib.client import Channel
+from quelware_core.entities.unit import UnitLabel, UnitStatus
+
+from quelware_client.core import AgentContainer, QuelwareClient
+from quelware_client.infra.instrument_agent_grpc import InstrumentAgentGrpc
+from quelware_client.infra.resource_agent_grpc import ResourceAgentGrpc
+from quelware_client.testing.session_agent_mock import SessionAgentMock
+from quelware_client.testing.system_configuration_agent_mock import (
+    SystemConfigurationAgentMock,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def create_standalone_client(
+    endpoint: str = "localhost", port: int = 50051, unit_label: str = "mock-unit"
+):
+    """Create a standalone client.
+
+    Note that the target worker server must be running with lock checking
+    disabled to use this client.
+    """
+    channel = Channel(endpoint, port)
+    target_unit = UnitLabel(unit_label)
+
+    conf_agent = SystemConfigurationAgentMock({target_unit: UnitStatus.ACTIVE})
+    session_agent = SessionAgentMock()
+
+    def resource_agent_factory(ul: UnitLabel):
+        return ResourceAgentGrpc(channel, metadata={"x-unit-label": str(ul)})
+
+    def instrument_agent_factory(ul: UnitLabel):
+        return InstrumentAgentGrpc(channel, metadata={"x-unit-label": str(ul)})
+
+    agent_container = AgentContainer()
+    agent_container.session = session_agent
+    agent_container.system_configuration = conf_agent
+
+    logger.warning("NOTE: Standalone client is for testing purposes.")
+
+    return QuelwareClient(
+        agent=agent_container,
+        resource_agent_factory=resource_agent_factory,
+        instrument_agent_factory=instrument_agent_factory,
+        close_handlers=[channel.close],
+    )

quelware_client/client/helpers/instrument_resolver/__init__.py

@@ -0,0 +1,58 @@
+import asyncio
+import logging
+
+from quelware_core.entities.instrument import InstrumentInfo
+from quelware_core.entities.resource import (
+    ResourceCategory,
+    ResourceId,
+)
+
+from quelware_client.core import QuelwareClient
+
+logger = logging.getLogger(__name__)
+
+
+class InstrumentResolver:
+    def __init__(self):
+        self._id_to_inst_info: dict[ResourceId, InstrumentInfo] = {}
+        self._alias_to_id: dict[str, ResourceId] = {}
+
+    async def refresh(self, client: QuelwareClient):
+        resource_infos = list(
+            rinfo
+            for rinfo in await client.list_resource_infos()
+            if rinfo.category is ResourceCategory.INSTRUMENT
+        )
+
+        coros = [client.get_instrument_info(rinfo.id) for rinfo in resource_infos]
+
+        inst_infos: list[InstrumentInfo] = await asyncio.gather(*coros)
+
+        new_id_to_inst_info = {}
+        new_alias_to_id = {}
+        for inst_info in inst_infos:
+            new_id_to_inst_info[inst_info.id] = inst_info
+            new_alias_to_id[inst_info.definition.alias] = inst_info.id
+
+        self._id_to_inst_info = new_id_to_inst_info
+        self._alias_to_id = new_alias_to_id
+        logger.info(f"{len(self._id_to_inst_info)} instruments has been registered.")
+
+    def find_inst_info_by_id(self, instrument_id: ResourceId) -> InstrumentInfo:
+        if instrument_id not in self._id_to_inst_info:
+            raise ValueError(f"Instrument with id '{instrument_id}' not found.")
+        return self._id_to_inst_info[instrument_id]
+
+    def find_inst_info_by_alias(self, alias: str) -> InstrumentInfo:
+        return self.find_inst_info_by_id(self._resolve_single(alias))
+
+    def _resolve_single(self, alias) -> ResourceId:
+        if alias not in self._alias_to_id:
+            raise ValueError(f"Instrument with alias '{alias}' not found.")
+        return self._alias_to_id[alias]
+
+    def resolve(self, aliases: list[str]) -> list[ResourceId]:
+        return list(self._resolve_single(alias) for alias in aliases)
+
+
+__all__ = ["InstrumentResolver"]

quelware_client/client/helpers/sequencer/__init__.py

@@ -0,0 +1,248 @@
+import logging
+import math
+from collections import defaultdict
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+from quelware_core.entities.directives import (
+    CaptureWindow,
+    SetFixedTimeline,
+    WaveformEvent,
+)
+from quelware_core.entities.waveform.sampled import IqArray, IqWaveform
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _SequencerEvent:
+    waveform_name: str
+    start_offset_ns: float
+    gain: float
+    phase_offset_deg: float
+
+
+@dataclass
+class _SequencerCaptureWindow:
+    name: str
+    start_offset_ns: float
+    length_ns: float
+
+
+@dataclass
+class _Waveform:
+    sampling_period_ns: float
+    iq_array: IqArray
+
+
+@dataclass
+class _AliasBinding:
+    sampling_period_fs: int
+    step_samples: int
+
+
+class Sequencer:
+    def __init__(
+        self, default_sampling_period_ns: float, enforce_sample_grid: bool = True
+    ):
+        self._waveform_library: dict[str, _Waveform] = {}
+        self._alias_to_events: dict[str, list[_SequencerEvent]] = defaultdict(list)
+        self._alias_to_capwin: dict[str, list[_SequencerCaptureWindow]] = defaultdict(
+            list
+        )
+
+        self._default_sampling_period_ns: float = default_sampling_period_ns
+        self._iterations: int = 1
+
+        self._bindings: dict[str, _AliasBinding] = {}
+        self._enforce_sample_grid: bool = enforce_sample_grid
+        self._length_ns: float = 0.0
+
+    def bind(self, alias: str, sampling_period_fs: int, step_samples: int):
+        self._bindings[alias] = _AliasBinding(
+            sampling_period_fs=sampling_period_fs,
+            step_samples=step_samples,
+        )
+
+    def _check_and_convert_to_samples(
+        self, alias: str, time_ns: float, name_for_log: str
+    ) -> int:
+        if alias not in self._bindings:
+            raise ValueError(
+                f"Alias '{alias}' is not bound to the sequencer. Call bind() first."
+            )
+
+        period_fs = self._bindings[alias].sampling_period_fs
+        time_fs = time_ns * 1e6
+        samples = time_fs / period_fs
+        rounded_samples = round(samples)
+
+        if abs(samples - rounded_samples) > 0.001:
+            msg = (
+                f"{name_for_log} ({time_ns} ns) is not a multiple of sampling period "
+                f"({period_fs} fs) for alias '{alias}'. Calculated samples: {samples}"
+            )
+            if self._enforce_sample_grid:
+                raise ValueError(msg)
+            else:
+                logger.warning(f"{msg}. Rounding to {rounded_samples} samples.")
+
+        return rounded_samples
+
+    def register_waveform(
+        self,
+        name: str,
+        waveform: npt.ArrayLike,
+        sampling_period_ns: float | None = None,
+    ):
+        if sampling_period_ns is None:
+            sampling_period_ns = self._default_sampling_period_ns
+        if np.any(np.abs(waveform) > 1):
+            raise ValueError("The amplitude must be in the range -1 to 1.")
+        self._waveform_library[name] = _Waveform(
+            sampling_period_ns=sampling_period_ns,
+            iq_array=np.array(waveform, dtype=complex),
+        )
+
+    def add_event(
+        self,
+        instrument_alias: str,
+        waveform_name: str,
+        start_offset_ns: float,
+        gain: float = 1.0,
+        phase_offset_deg: float = 0.0,
+    ):
+        if waveform_name not in self._waveform_library:
+            raise ValueError(f"waveform '{waveform_name}' is not registered.")
+
+        self._check_and_convert_to_samples(
+            instrument_alias, start_offset_ns, "Event start_offset_ns"
+        )
+
+        event = _SequencerEvent(
+            waveform_name=waveform_name,
+            start_offset_ns=start_offset_ns,
+            gain=gain,
+            phase_offset_deg=phase_offset_deg,
+        )
+        self._alias_to_events[instrument_alias].append(event)
+
+        waveform = self._waveform_library[waveform_name]
+        end_at_ns = (
+            start_offset_ns + len(waveform.iq_array) * waveform.sampling_period_ns
+        )
+        self._length_ns = max(self._length_ns, end_at_ns)
+
+    def add_capture_window(
+        self,
+        instrument_alias: str,
+        window_name: str,
+        start_offset_ns: float,
+        length_ns: float,
+    ):
+        self._check_and_convert_to_samples(
+            instrument_alias,
+            start_offset_ns,
+            f"Capture window '{window_name}' start_offset_ns",
+        )
+        self._check_and_convert_to_samples(
+            instrument_alias, length_ns, f"Capture window '{window_name}' length_ns"
+        )
+
+        capwin = _SequencerCaptureWindow(
+            name=window_name,
+            start_offset_ns=start_offset_ns,
+            length_ns=length_ns,
+        )
+        self._alias_to_capwin[instrument_alias].append(capwin)
+        end_at_ns = start_offset_ns + length_ns
+        self._length_ns = max(self._length_ns, end_at_ns)
+
+    def extend_length_ns(self, additional_ns: float):
+        self._length_ns += additional_ns
+
+    def set_iterations(self, iterations: int):
+        self._iterations = iterations
+
+    @property
+    def aligned_length_fs(self) -> int:
+        if not self._bindings:
+            return math.ceil(self._length_ns * 1e6)
+
+        lcm_step_fs = 1
+        for b in self._bindings.values():
+            step_fs = b.sampling_period_fs * b.step_samples
+            lcm_step_fs = math.lcm(lcm_step_fs, step_fs)
+
+        length_fs = math.ceil(self._length_ns * 1e6)
+        remainder = length_fs % lcm_step_fs
+        if remainder != 0:
+            length_fs += lcm_step_fs - remainder
+        return length_fs
+
+    def export_set_fixed_timeline_directive(
+        self, instrument_alias: str
+    ) -> SetFixedTimeline:
+        if instrument_alias not in self._bindings:
+            raise ValueError(f"Alias '{instrument_alias}' is not bound.")
+
+        sampling_period_fs = self._bindings[instrument_alias].sampling_period_fs
+
+        name_to_index: dict[str, int] = {}
+        counter = 0
+        local_library: list[IqWaveform] = []
+        local_events: list[WaveformEvent] = []
+
+        for event in self._alias_to_events[instrument_alias]:
+            if event.waveform_name in name_to_index:
+                index = name_to_index[event.waveform_name]
+            else:
+                waveform = self._waveform_library[event.waveform_name]
+                local_library.append(
+                    IqWaveform(
+                        sampling_period_fs=round(waveform.sampling_period_ns * 1e6),
+                        iq_array=waveform.iq_array,
+                    )
+                )
+                index = counter
+                name_to_index[event.waveform_name] = index
+                counter += 1
+
+            start_offset_samples = self._check_and_convert_to_samples(
+                instrument_alias, event.start_offset_ns, "Export event"
+            )
+            local_events.append(
+                WaveformEvent(
+                    waveform_index=index,
+                    start_offset_samples=start_offset_samples,
+                    gain=event.gain,
+                    phase_offset_deg=event.phase_offset_deg,
+                )
+            )
+
+        local_capwins: list[CaptureWindow] = []
+        for capwin in self._alias_to_capwin[instrument_alias]:
+            local_capwins.append(
+                CaptureWindow(
+                    name=capwin.name,
+                    start_offset_samples=self._check_and_convert_to_samples(
+                        instrument_alias, capwin.start_offset_ns, "Export capwin start"
+                    ),
+                    length_samples=self._check_and_convert_to_samples(
+                        instrument_alias, capwin.length_ns, "Export capwin length"
+                    ),
+                )
+            )
+
+        length_sample = (
+            self.aligned_length_fs + sampling_period_fs - 1
+        ) // sampling_period_fs
+
+        return SetFixedTimeline(
+            waveform_library=local_library,
+            events=local_events,
+            capture_windows=local_capwins,
+            length=length_sample,
+            iterations=self._iterations,
+        )

quelware_client/core/__init__.py

@@ -0,0 +1,5 @@
+from ._agent_container import AgentContainer
+from ._client import AgentFactory, QuelwareClient
+from ._session import Session
+
+__all__ = ["AgentContainer", "AgentFactory", "QuelwareClient", "Session"]

quelware_client/core/_agent_container.py

@@ -0,0 +1,72 @@
+import logging
+
+from quelware_core.entities.unit import UnitLabel
+
+from quelware_client.core.interfaces.instrument_agent import InstrumentAgent
+from quelware_client.core.interfaces.resource_agent import ResourceAgent
+from quelware_client.core.interfaces.session_agent import SessionAgent
+from quelware_client.core.interfaces.system_configuration_agent import (
+    SystemConfigurationAgent,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AgentContainer:
+    def __init__(self):
+        self._session: SessionAgent | None = None
+        self._conf: SystemConfigurationAgent | None = None
+        self._rsrc: dict[UnitLabel, ResourceAgent] = {}
+        self._inst: dict[UnitLabel, InstrumentAgent] = {}
+
+    def update_resource_agent(self, unit_label: UnitLabel, agent: ResourceAgent | None):
+        if agent is None:
+            if unit_label in self._rsrc:
+                self._rsrc[unit_label]
+            else:
+                logger.warning(f"ResourceAgent for '{unit_label}' not found. Continue.")
+        else:
+            self._rsrc[unit_label] = agent
+
+    def update_instrument_agent(
+        self, unit_label: UnitLabel, agent: InstrumentAgent | None
+    ):
+        if agent is None:
+            if unit_label in self._rsrc:
+                self._inst[unit_label]
+            else:
+                logger.warning(
+                    f"InstrumentAgent for '{unit_label}' not found. Continue."
+                )
+        else:
+            self._inst[unit_label] = agent
+
+    @property
+    def session(self) -> SessionAgent:
+        if self._session is None:
+            raise ValueError("SessionAgent not set.")
+        return self._session
+
+    @session.setter
+    def session(self, val: SessionAgent):
+        self._session = val
+
+    @property
+    def system_configuration(self) -> SystemConfigurationAgent:
+        if self._conf is None:
+            raise ValueError("ConfigurationAgent not set.")
+        return self._conf
+
+    @system_configuration.setter
+    def system_configuration(self, val: SystemConfigurationAgent):
+        self._conf = val
+
+    def resource(self, unit_label: UnitLabel) -> ResourceAgent:
+        if unit_label not in self._rsrc:
+            raise ValueError(f"ResouceAgent for '{unit_label}' not set.")
+        return self._rsrc[unit_label]
+
+    def instrument(self, unit_label: UnitLabel) -> InstrumentAgent:
+        if unit_label not in self._inst:
+            raise ValueError(f"InstrumentAgent for '{unit_label}' not set.")
+        return self._inst[unit_label]

quelware_client/core/_client.py

@@ -0,0 +1,99 @@
+import asyncio
+from collections.abc import Callable, Collection
+from typing import TypeAlias, TypeVar, cast
+
+from quelware_core.entities.instrument import InstrumentInfo
+from quelware_core.entities.port import PortInfo
+from quelware_core.entities.resource import (
+    ResourceId,
+    ResourceInfo,
+    extract_unit_label,
+)
+from quelware_core.entities.unit import UnitLabel
+
+from quelware_client.core._agent_container import AgentContainer
+from quelware_client.core.interfaces.instrument_agent import InstrumentAgent
+
+from ._session import Session
+from .interfaces.resource_agent import ResourceAgent
+
+A = TypeVar("A")
+AgentFactory: TypeAlias = Callable[[UnitLabel], A]
+_ResourceId: TypeAlias = ResourceId | str
+
+
+class QuelwareClient:
+    def __init__(
+        self,
+        agent: AgentContainer,
+        resource_agent_factory: AgentFactory[ResourceAgent] | None = None,
+        instrument_agent_factory: AgentFactory[InstrumentAgent] | None = None,
+        close_handlers: list[Callable[[], None]] | None = None,
+    ):
+        self._agent = agent
+        self._rsrc_agent_factory = resource_agent_factory
+        self._inst_agent_factory = instrument_agent_factory
+        self._unit_labels: list[UnitLabel] = []
+        self._close_handlers = close_handlers or []
+
+    @property
+    def agent(self) -> AgentContainer:
+        return self._agent
+
+    async def start(self):
+        await self.initialize()
+
+    async def stop(self):
+        for handler in self._close_handlers:
+            handler()
+
+    async def __aenter__(self):
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.stop()
+
+    async def initialize(self):
+        self._unit_labels = await self._agent.system_configuration.list_units()
+        for ul in self._unit_labels:
+            if self._rsrc_agent_factory:
+                self._agent.update_resource_agent(ul, self._rsrc_agent_factory(ul))
+            if self._inst_agent_factory:
+                self._agent.update_instrument_agent(ul, self._inst_agent_factory(ul))
+
+    def list_unit_labels(self) -> list[UnitLabel]:
+        return list(self._unit_labels)
+
+    def create_session(
+        self,
+        resource_ids: Collection[_ResourceId],
+        ttl_ms: int = 4000,
+        tentative_ttl_ms: int = 1000,
+    ) -> Session:
+        return Session(
+            resource_ids=cast(list[ResourceId], resource_ids),
+            agent=self._agent,
+            ttl_ms=ttl_ms,
+            tentative_ttl_ms=tentative_ttl_ms,
+        )
+
+    async def list_resource_infos(self) -> list[ResourceInfo]:
+        coros = [
+            self._agent.resource(ul).list_resource_infos() for ul in self._unit_labels
+        ]
+        results = await asyncio.gather(*coros)
+        rsrc_infos: list[ResourceInfo] = []
+        for res in results:
+            rsrc_infos.extend(res)
+        return rsrc_infos
+
+    async def get_port_info(self, port_id: ResourceId) -> PortInfo:
+        unit_label = extract_unit_label(port_id)
+        port = await self._agent.resource(unit_label).get_port_info(port_id)
+        return port
+
+    async def get_instrument_info(self, instrument_id: ResourceId) -> InstrumentInfo:
+        unit_label = extract_unit_label(instrument_id)
+        inst = await self._agent.resource(unit_label).get_instrument_info(instrument_id)
+        return inst