diff --git a/sequence/resource_management/reservation_rule_registry.py b/sequence/resource_management/reservation_rule_registry.py new file mode 100644 index 000000000..c30de0bf6 --- /dev/null +++ b/sequence/resource_management/reservation_rule_registry.py @@ -0,0 +1,303 @@ +"""Utilities for building reservation-generated resource-management rules. + +This module defines the fixed reservation rule slots used by the default +reservation rule generator. Users may replace or disable the builders for +these existing slots, but the reservation rule order and creation predicates +remain static. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from inspect import signature +from typing import TYPE_CHECKING, Final +from collections.abc import Callable, Iterable + +from .action_condition_set import ( + eg_rule_action_await, + eg_rule_action_request, + eg_rule_condition, + ep_rule_action_await, + ep_rule_action_request, + ep_rule_condition_await, + ep_rule_condition_request, + es_rule_action_A, + es_rule_action_B, + es_rule_condition_A, + es_rule_condition_B, + es_rule_condition_B_end, +) +from .rule_manager import Rule + +if TYPE_CHECKING: + from ..network_management.reservation import Reservation + from ..topology.node import QuantumRouter + + +EG_AWAIT: Final[str] = "eg_await" +EG_REQUEST: Final[str] = "eg_request" +EP_REQUEST: Final[str] = "ep_request" +EP_AWAIT: Final[str] = "ep_await" +ES_B_END: Final[str] = "es_b_end" +ES_A: Final[str] = "es_a" +ES_B: Final[str] = "es_b" + +RESERVATION_RULE_NAMES: Final[tuple[str, ...]] = ( + EG_AWAIT, + EG_REQUEST, + EP_REQUEST, + EP_AWAIT, + ES_B_END, + ES_A, + ES_B, +) + + +@dataclass(frozen=True) +class ReservationRuleContext: + """Inputs needed to build reservation-generated resource-management rules.""" + + owner: QuantumRouter + path: list[str] + reservation: Reservation + memory_indices: list[int] + index: int + + +ReservationRulePredicate = Callable[[ReservationRuleContext], bool] +ReservationRuleBuilder = Callable[[ReservationRuleContext], Rule] + + +@dataclass(frozen=True) +class ReservationRuleSpec: + """Static specification for one built-in reservation rule slot.""" + + name: str + predicate: ReservationRulePredicate + builder: ReservationRuleBuilder + + +def _applies_eg_await(context: ReservationRuleContext) -> bool: + return context.index > 0 + + +def _applies_eg_request(context: ReservationRuleContext) -> bool: + return context.index < len(context.path) - 1 + + +def _applies_ep_request(context: ReservationRuleContext) -> bool: + return context.index > 0 + + +def _applies_ep_await(context: ReservationRuleContext) -> bool: + return context.index < len(context.path) - 1 + + +def _applies_es_b_end(context: ReservationRuleContext) -> bool: + return context.index == 0 or context.index == len(context.path) - 1 + + +def _applies_es_middle(context: ReservationRuleContext) -> bool: + return 0 < context.index < len(context.path) - 1 + + +def _build_eg_await_rule(context: ReservationRuleContext) -> Rule: + condition_args = {"memory_indices": context.memory_indices[:context.reservation.memory_size]} + action_args = { + "mid": context.owner.map_to_middle_node[context.path[context.index - 1]], + "path": context.path, + "index": context.index, + } + return Rule(10, eg_rule_action_await, eg_rule_condition, action_args, condition_args) + + +def _build_eg_request_rule(context: ReservationRuleContext) -> Rule: + if context.index == 0: + condition_args = {"memory_indices": context.memory_indices[:context.reservation.memory_size]} + else: + condition_args = {"memory_indices": context.memory_indices[context.reservation.memory_size:]} + + action_args = { + "mid": context.owner.map_to_middle_node[context.path[context.index + 1]], + "path": context.path, + "index": context.index, + "name": context.owner.name, + "reservation": context.reservation, + } + return Rule(10, eg_rule_action_request, eg_rule_condition, action_args, condition_args) + + +def _build_ep_request_rule(context: ReservationRuleContext) -> Rule: + condition_args = { + "memory_indices": context.memory_indices[:context.reservation.memory_size], + "reservation": context.reservation, + "purification_mode": context.reservation.purification_mode, + } + return Rule(10, ep_rule_action_request, ep_rule_condition_request, {}, condition_args) + + +def _build_ep_await_rule(context: ReservationRuleContext) -> Rule: + if context.index == 0: + condition_args = { + "memory_indices": context.memory_indices, + "fidelity": context.reservation.fidelity, + "purification_mode": context.reservation.purification_mode, + } + else: + condition_args = { + "memory_indices": context.memory_indices[context.reservation.memory_size:], + "fidelity": context.reservation.fidelity, + "purification_mode": context.reservation.purification_mode, + } + + return Rule(10, ep_rule_action_await, ep_rule_condition_await, {}, condition_args) + + +def _build_es_b_end_rule(context: ReservationRuleContext) -> Rule: + if context.index == 0: + target_remote = context.path[-1] + else: + target_remote = context.path[0] + + condition_args = { + "memory_indices": context.memory_indices, + "target_remote": target_remote, + "fidelity": context.reservation.fidelity, + } + return Rule(10, es_rule_action_B, es_rule_condition_B_end, {}, condition_args) + + +def _get_swapping_neighbors(owner_name: str, path: list[str]) -> tuple[str, str]: + """Return the left and right neighbors used for middle-node swapping rules.""" + reduced_path = path[:] + while reduced_path.index(owner_name) % 2 == 0: + new_path = [] + for i, node_name in enumerate(reduced_path): + if i % 2 == 0 or i == len(reduced_path) - 1: + new_path.append(node_name) + reduced_path = new_path + + reduced_index = reduced_path.index(owner_name) + return reduced_path[reduced_index - 1], reduced_path[reduced_index + 1] + + +def _build_es_a_rule(context: ReservationRuleContext) -> Rule: + left, right = _get_swapping_neighbors(context.owner.name, context.path) + condition_args = { + "memory_indices": context.memory_indices, + "left": left, + "right": right, + "fidelity": context.reservation.fidelity, + } + action_args = { + "swapping_success_prob": context.owner.swapping_success_prob, + "swapping_degradation": context.owner.swapping_degradation, + } + return Rule(10, es_rule_action_A, es_rule_condition_A, action_args, condition_args) + + +def _build_es_b_rule(context: ReservationRuleContext) -> Rule: + left, right = _get_swapping_neighbors(context.owner.name, context.path) + condition_args = { + "memory_indices": context.memory_indices, + "left": left, + "right": right, + "fidelity": context.reservation.fidelity, + } + return Rule(10, es_rule_action_B, es_rule_condition_B, {}, condition_args) + + +DEFAULT_RESERVATION_RULE_SPECS: Final[tuple[ReservationRuleSpec, ...]] = ( + ReservationRuleSpec(EG_AWAIT, _applies_eg_await, _build_eg_await_rule), + ReservationRuleSpec(EG_REQUEST, _applies_eg_request, _build_eg_request_rule), + ReservationRuleSpec(EP_REQUEST, _applies_ep_request, _build_ep_request_rule), + ReservationRuleSpec(EP_AWAIT, _applies_ep_await, _build_ep_await_rule), + ReservationRuleSpec(ES_B_END, _applies_es_b_end, _build_es_b_end_rule), + ReservationRuleSpec(ES_A, _applies_es_middle, _build_es_a_rule), + ReservationRuleSpec(ES_B, _applies_es_middle, _build_es_b_rule), +) + + +class ReservationRuleRegistry: + """Registry for replacing or disabling existing reservation rule builders.""" + + def __init__(self, specs: Iterable[ReservationRuleSpec] | None = None) -> None: + if specs is None: + specs = DEFAULT_RESERVATION_RULE_SPECS + self._builders: dict[str, ReservationRuleBuilder] = {} + for spec in specs: + self._validate_name(spec.name) + self._validate_builder(spec.builder) + if spec.name in self._builders: + msg = f"Duplicate reservation rule name: {spec.name}" + raise ValueError(msg) + self._builders[spec.name] = spec.builder + + def replace(self, name: str, builder: ReservationRuleBuilder) -> None: + """Replace the builder for an existing reservation rule slot.""" + self._validate_name(name) + self._validate_builder(builder) + self._builders[name] = builder + + def disable(self, name: str) -> None: + """Disable an existing reservation rule slot.""" + self._validate_name(name) + self._builders.pop(name, None) + + def get(self, name: str) -> ReservationRuleBuilder | None: + """Return the builder for a rule slot, or None if it is disabled.""" + self._validate_name(name) + return self._builders.get(name) + + def build(self, name: str, context: ReservationRuleContext) -> Rule | None: + """Build the named rule for the supplied context. + + Returns None when the rule slot is disabled. + """ + builder = self.get(name) + if builder is None: + return None + + rule = builder(context) + if not isinstance(rule, Rule): + msg = f"Reservation rule builder for {name!r} must return a Rule" + raise TypeError(msg) + return rule + + @staticmethod + def _validate_name(name: str) -> None: + if name not in RESERVATION_RULE_NAMES: + msg = f"Unknown reservation rule name: {name}" + raise ValueError(msg) + + @staticmethod + def _validate_builder(builder: ReservationRuleBuilder) -> None: + if not callable(builder): + msg = "Reservation rule builders must be callable" + raise TypeError(msg) + + try: + signature(builder).bind(object()) + except (TypeError, ValueError) as exc: + msg = "Reservation rule builders must accept a ReservationRuleContext argument" + raise TypeError(msg) from exc + + +class ReservationRuleGenerator: + """Rule generator for reservation-based resource-management rules.""" + + def __init__(self) -> None: + self.registry = ReservationRuleRegistry() + + def create_rules(self, owner, path, reservation, memory_indices, index) -> list[Rule]: + """Create reservation rules in SeQUeNCe's static rule-slot order.""" + context = ReservationRuleContext(owner, path, reservation, memory_indices, index) + rules: list[Rule] = [] + + for spec in DEFAULT_RESERVATION_RULE_SPECS: + if spec.predicate(context): + rule = self.registry.build(spec.name, context) + if rule is not None: + rules.append(rule) + + return rules diff --git a/sequence/resource_management/resource_manager.py b/sequence/resource_management/resource_manager.py index 62c552f3d..ca6d9707a 100644 --- a/sequence/resource_management/resource_manager.py +++ b/sequence/resource_management/resource_manager.py @@ -10,19 +10,9 @@ from typing import TYPE_CHECKING from collections.abc import Callable -from .action_condition_set import ( - eg_rule_action_await, - eg_rule_action_request, - eg_rule_condition, - ep_rule_action_await, - ep_rule_action_request, - ep_rule_condition_await, - ep_rule_condition_request, - es_rule_action_A, - es_rule_action_B, - es_rule_condition_A, - es_rule_condition_B, - es_rule_condition_B_end, +from .reservation_rule_registry import ( + ReservationRuleGenerator, + ReservationRuleRegistry, ) from ..kernel.event import Event from ..kernel.process import Process @@ -148,10 +138,15 @@ def __init__(self, owner: QuantumRouter, memory_array_name: str): self.pending_protocols = [] # Protocols that are requesting remote resource self.waiting_protocols = [] # Protocols that are waiting request from remote resource self.memory_to_protocol_map = {} + self.rule_generator = ReservationRuleGenerator() def init(self): pass + def get_reservation_rule_registry(self) -> ReservationRuleRegistry: + """Return the registry used to replace or disable reservation rule builders.""" + return self.rule_generator.registry + def generate_load_rules(self, path: list[str], reservation: Reservation, timecards: list[MemoryTimeCard], memory_array_name: str): """Generate and load rules for a given reservation. @@ -161,7 +156,6 @@ def generate_load_rules(self, path: list[str], reservation: Reservation, timecar timecards (list[MemoryTimeCard]): timecards involved in the reservation at this node. memory_array_name (str): name of memory array component to use for rule conditions and actions at this node. """ - rules = [] memory_indices = [] for card in timecards: if reservation in card.reservations: @@ -169,85 +163,13 @@ def generate_load_rules(self, path: list[str], reservation: Reservation, timecar log.logger.debug(f'Memory indices for reservation {reservation.identity} on node {self.owner.name}: {memory_indices}') index: int = path.index(self.owner.name) - - # Create Rules - # 1. create rules for entanglement generation - if index > 0: - condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} - action_args = {"mid": self.owner.map_to_middle_node[path[index - 1]], - "path": path, "index": index} - rule = Rule(10, eg_rule_action_await, eg_rule_condition, action_args, condition_args) - rules.append(rule) - - if index < len(path) - 1: - if index == 0: - condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} - else: - condition_args = {"memory_indices": memory_indices[reservation.memory_size:]} - - action_args = {"mid": self.owner.map_to_middle_node[path[index + 1]], - "path": path, "index": index, "name": self.owner.name, "reservation": reservation} - rule = Rule(10, eg_rule_action_request, eg_rule_condition, action_args, condition_args) - rules.append(rule) - - # 2. create rules for entanglement purification - if index > 0: - condition_args = {"memory_indices": memory_indices[:reservation.memory_size], "reservation": reservation, - "purification_mode": reservation.purification_mode} - action_args = {} - rule = Rule(10, ep_rule_action_request, ep_rule_condition_request, action_args, condition_args) - rules.append(rule) - - if index < len(path) - 1: - if index == 0: - condition_args = {"memory_indices": memory_indices, "fidelity": reservation.fidelity, - "purification_mode": reservation.purification_mode} - else: - condition_args = {"memory_indices": memory_indices[reservation.memory_size:], - "fidelity": reservation.fidelity, - "purification_mode": reservation.purification_mode} - - action_args = {} - rule = Rule(10, ep_rule_action_await, ep_rule_condition_await, action_args, condition_args) - rules.append(rule) - - # 3. create rules for entanglement swapping - if index == 0: - condition_args = {"memory_indices": memory_indices, "target_remote": path[-1], - "fidelity": reservation.fidelity} - action_args = {} - rule = Rule(10, es_rule_action_B, es_rule_condition_B_end, action_args, condition_args) - rules.append(rule) - - elif index == len(path) - 1: - action_args = {} - condition_args = {"memory_indices": memory_indices, "target_remote": path[0], - "fidelity": reservation.fidelity} - rule = Rule(10, es_rule_action_B, es_rule_condition_B_end, action_args, condition_args) - rules.append(rule) - - else: - _path = path[:] - while _path.index(self.owner.name) % 2 == 0: - new_path = [] - for i, n in enumerate(_path): - if i % 2 == 0 or i == len(_path) - 1: - new_path.append(n) - _path = new_path - _index = _path.index(self.owner.name) - left, right = _path[_index - 1], _path[_index + 1] - - condition_args = {"memory_indices": memory_indices, "left": left, "right": right, - "fidelity": reservation.fidelity} - action_args = {"swapping_success_prob": self.owner.swapping_success_prob, - "swapping_degradation": self.owner.swapping_degradation} - rule = Rule(10, es_rule_action_A, es_rule_condition_A, action_args, condition_args) - rules.append(rule) - - action_args = {} - rule = Rule(10, es_rule_action_B, es_rule_condition_B, action_args, condition_args) - rules.append(rule) - + rules = self.rule_generator.create_rules( + self.owner, + path, + reservation, + memory_indices, + index, + ) for rule in rules: rule.set_reservation(reservation) diff --git a/tests/resource_management/test_reservation_rule_registry.py b/tests/resource_management/test_reservation_rule_registry.py new file mode 100644 index 000000000..6c90cedbf --- /dev/null +++ b/tests/resource_management/test_reservation_rule_registry.py @@ -0,0 +1,171 @@ +import pytest + +from sequence.resource_management.action_condition_set import ( + eg_rule_action_request, + ep_rule_action_await, + es_rule_action_A, +) +from sequence.resource_management.reservation_rule_registry import ( + DEFAULT_RESERVATION_RULE_SPECS, + EG_AWAIT, + EG_REQUEST, + EP_AWAIT, + EP_REQUEST, + ES_A, + ES_B, + ES_B_END, + ReservationRuleContext, + ReservationRuleGenerator, + ReservationRuleRegistry, +) +from sequence.resource_management.rule_manager import Rule + + +class Owner: + name = "node1" + map_to_middle_node = {"node0": "middle_left", "node2": "middle_right"} + swapping_success_prob = 0.75 + swapping_degradation = 0.95 + + +class Reservation: + memory_size = 1 + fidelity = 0.9 + purification_mode = "BBPSSW" + + +def _custom_rule_builder(context: ReservationRuleContext) -> Rule: + return Rule(10, eg_rule_action_request, lambda *_args: [], {"custom": context.index}, {}) + + +def test_registry_can_disable_individual_rule_builder(): + registry = ReservationRuleRegistry() + + assert registry.get(EP_REQUEST) is not None + + registry.disable(EP_REQUEST) + + assert registry.get(EP_REQUEST) is None + assert registry.get(EG_REQUEST) is not None + assert registry.get(EP_AWAIT) is not None + + +def test_registry_can_replace_individual_rule_builder(): + registry = ReservationRuleRegistry() + + registry.replace(EP_REQUEST, _custom_rule_builder) + + assert registry.get(EP_REQUEST) is _custom_rule_builder + assert registry.get(EG_REQUEST) is not _custom_rule_builder + + +def test_registry_rejects_unknown_rule_name(): + registry = ReservationRuleRegistry() + + with pytest.raises(ValueError): + registry.replace("custom_phase", _custom_rule_builder) + + +def test_registry_rejects_builder_with_wrong_signature(): + registry = ReservationRuleRegistry() + + def bad_builder(): + return Rule(10, eg_rule_action_request, lambda *_args: [], {}, {}) + + with pytest.raises(TypeError): + registry.replace(EP_REQUEST, bad_builder) + + +def test_generator_uses_replacement_builder(): + generator = ReservationRuleGenerator() + generator.registry.replace(EG_REQUEST, _custom_rule_builder) + + rules = generator.create_rules( + Owner(), + ["node1", "node2"], + Reservation(), + [0], + 0, + ) + + assert all(isinstance(rule, Rule) for rule in rules) + assert any(rule.action_args == {"custom": 0} for rule in rules) + assert not any(rule.action is eg_rule_action_request and "custom" not in rule.action_args for rule in rules) + + +def test_generator_uses_registry_to_disable_rule(): + generator = ReservationRuleGenerator() + registry = generator.registry + + rules = generator.create_rules( + Owner(), + ["node1", "node2"], + Reservation(), + [0], + 0, + ) + + assert all(isinstance(rule, Rule) for rule in rules) + assert any(rule.action is eg_rule_action_request for rule in rules) + assert any(rule.action is ep_rule_action_await for rule in rules) + + registry.disable(EP_AWAIT) + + rules = generator.create_rules( + Owner(), + ["node1", "node2"], + Reservation(), + [0], + 0, + ) + + assert any(rule.action is eg_rule_action_request for rule in rules) + assert not any(rule.action is ep_rule_action_await for rule in rules) + + +def test_generator_uses_static_default_rule_spec_order(): + generator = ReservationRuleGenerator() + + def marker_builder(name): + def build(_context: ReservationRuleContext) -> Rule: + return Rule(10, eg_rule_action_request, lambda *_args: [], {"slot": name}, {}) + + return build + + for rule_name in (EG_AWAIT, EG_REQUEST, EP_REQUEST, EP_AWAIT, ES_B_END, ES_A, ES_B): + generator.registry.replace(rule_name, marker_builder(rule_name)) + + rules = generator.create_rules( + Owner(), + ["node0", "node1", "node2"], + Reservation(), + [0, 1], + 1, + ) + + context = ReservationRuleContext(Owner(), ["node0", "node1", "node2"], Reservation(), [0, 1], 1) + expected_names = [spec.name for spec in DEFAULT_RESERVATION_RULE_SPECS if spec.predicate(context)] + actual_names = [rule.action_args["slot"] for rule in rules] + + assert actual_names == expected_names + + +def test_default_es_a_rule_preserves_swapping_action_args(): + generator = ReservationRuleGenerator() + + rules = generator.create_rules( + Owner(), + ["node0", "node1", "node2"], + Reservation(), + [0, 1], + 1, + ) + + es_a_rules = [rule for rule in rules if rule.action is es_rule_action_A] + assert len(es_a_rules) == 1 + assert es_a_rules[0].action_args == { + "swapping_success_prob": Owner.swapping_success_prob, + "swapping_degradation": Owner.swapping_degradation, + } + assert es_a_rules[0].condition_args["left"] == "node0" + assert es_a_rules[0].condition_args["right"] == "node2"