diff --git a/sequence/resource_management/resource_manager.py b/sequence/resource_management/resource_manager.py index 17d22ab73..c04503c5c 100644 --- a/sequence/resource_management/resource_manager.py +++ b/sequence/resource_management/resource_manager.py @@ -10,20 +10,6 @@ from typing import TYPE_CHECKING, Optional 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 ..kernel.event import Event from ..kernel.process import Process @@ -31,6 +17,7 @@ from ..components.memory import Memory from ..topology.node import QuantumRouter from .rule_manager import Rule, Arguments +from .rule_generation import DefaultReservationRuleGenerator from ..entanglement_management.entanglement_protocol import EntanglementProtocol from ..message import Message @@ -141,10 +128,21 @@ def __init__(self, owner: QuantumRouter, memory_array_name: str): self.memory_manager.set_resource_manager(self) self.rule_manager = RuleManager() self.rule_manager.set_resource_manager(self) + self.rule_generator = DefaultReservationRuleGenerator() 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 = {} + def set_rule_generator(self, rule_generator: DefaultReservationRuleGenerator) -> None: + """Set the generator used to create reservation rules. + + Args: + rule_generator (DefaultReservationRuleGenerator): generator used to create + reservation-based entanglement management rules. + """ + + self.rule_generator = rule_generator + 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. @@ -154,7 +152,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: @@ -162,82 +159,13 @@ def generate_load_rules(self, path: list[str], reservation: Reservation, timecar 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 = {} - 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/sequence/resource_management/rule_generation.py b/sequence/resource_management/rule_generation.py new file mode 100644 index 000000000..e370f7d93 --- /dev/null +++ b/sequence/resource_management/rule_generation.py @@ -0,0 +1,198 @@ +"""Default rule generation for reservation-based entanglement management. + +This module defines the default generator used by the resource manager to create +entanglement generation, purification, and swapping rules for a reservation. + +The default behavior matches ResourceManager.generate_load_rules, while allowing +subclasses to override parts of reservation rule creation. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +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 + + +class DefaultReservationRuleGenerator: + """Default generator for reservation-created entanglement management rules. + + The methods in this class preserve the original rule-construction behavior + of ResourceManager.generate_load_rules while allowing subclasses to override + generation, purification, or swapping rule creation independently. + """ + + def create_generation_rules( + self, + owner: "QuantumRouter", + path: list[str], + reservation: "Reservation", + memory_indices: list[int], + index: int, + ) -> list[Rule]: + """Create entanglement generation rules for this node.""" + + rules = [] + + if index > 0: + condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} + action_args = { + "mid": 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": owner.map_to_middle_node[path[index + 1]], + "path": path, + "index": index, + "name": owner.name, + "reservation": reservation, + } + rule = Rule(10, eg_rule_action_request, eg_rule_condition, action_args, condition_args) + rules.append(rule) + + return rules + + def create_purification_rules( + self, + _owner: "QuantumRouter", + path: list[str], + reservation: "Reservation", + memory_indices: list[int], + index: int, + ) -> list[Rule]: + """Create entanglement purification rules for this node.""" + + rules = [] + + 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) + return rules + + def create_swapping_rules( + self, + owner: "QuantumRouter", + path: list[str], + reservation: "Reservation", + memory_indices: list[int], + index: int, + ) -> list[Rule]: + """Create entanglement swapping rules for this node.""" + + rules = [] + + 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(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(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 = {} + 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) + + return rules + + def create_rules( + self, + owner: "QuantumRouter", + path: list[str], + reservation: "Reservation", + memory_indices: list[int], + index: int, + ) -> list[Rule]: + """Create all default rules for this reservation at this node.""" + + rules = [] + rules.extend(self.create_generation_rules(owner, path, reservation, memory_indices, index)) + rules.extend(self.create_purification_rules(owner, path, reservation, memory_indices, index)) + rules.extend(self.create_swapping_rules(owner, path, reservation, memory_indices, index)) + return rules diff --git a/tests/resource_management/test_rule_generation.py b/tests/resource_management/test_rule_generation.py new file mode 100644 index 000000000..d646b923d --- /dev/null +++ b/tests/resource_management/test_rule_generation.py @@ -0,0 +1,166 @@ +from sequence.kernel.timeline import Timeline +from sequence.network_management.reservation import Reservation +from sequence.resource_management.rule_generation import DefaultReservationRuleGenerator +from sequence.resource_management.action_condition_set import ( + eg_rule_action_request, + eg_rule_action_await, + ep_rule_action_request, + ep_rule_action_await, + es_rule_action_B, +) +from sequence.topology.node import Node +from sequence.components.memory import MemoryArray +from sequence.resource_management.resource_manager import ResourceManager + + +class FakeOwner(Node): + def __init__(self, name, tl): + super().__init__(name, tl) + self.map_to_middle_node = { + "node2": "mid_node1", + "node1": "mid_node1", + } + self.memory_array_name = name + ".MemoryArray" + memory_array = MemoryArray(self.memory_array_name, tl) + self.add_component(memory_array) + self.resource_manager = ResourceManager(self, self.memory_array_name) + + def get_idle_memory(self, info): + pass + + +def test_default_reservation_rule_generator_creates_direct_path_rules(): + tl = Timeline() + owner = FakeOwner("node1", tl) + generator = DefaultReservationRuleGenerator() + + reservation = Reservation( + "node1", + "node2", + start_time=0, + end_time=1e12, + memory_size=1, + fidelity=0.9, + entanglement_number=1, + identity=0, + ) + + rules = generator.create_rules( + owner=owner, + path=["node1", "node2"], + reservation=reservation, + memory_indices=[0], + index=0, + ) + + actions = [rule.action for rule in rules] + + assert eg_rule_action_request in actions + assert eg_rule_action_await not in actions + assert ep_rule_action_await in actions + assert ep_rule_action_request not in actions + assert es_rule_action_B in actions + + +class TrackingRuleGenerator(DefaultReservationRuleGenerator): + def __init__(self): + self.called = False + self.call_args = None + + def create_rules(self, owner, path, reservation, memory_indices, index): + self.called = True + self.call_args = { + "owner": owner, + "path": path, + "reservation": reservation, + "memory_indices": memory_indices, + "index": index, + } + return [] + + +class FakeTimeCard: + def __init__(self, memory_index, reservations): + self.memory_index = memory_index + self.reservations = reservations + + +def test_resource_manager_uses_custom_rule_generator(): + tl = Timeline() + owner = FakeOwner("node1", tl) + generator = TrackingRuleGenerator() + + reservation = Reservation( + "node1", + "node2", + start_time=0, + end_time=1e12, + memory_size=1, + fidelity=0.9, + entanglement_number=1, + identity=0, + ) + + timecards = [ + FakeTimeCard(0, [reservation]), + FakeTimeCard(1, []), + ] + + owner.resource_manager.set_rule_generator(generator) + owner.resource_manager.generate_load_rules( + path=["node1", "node2"], + reservation=reservation, + timecards=timecards, + memory_array_name=owner.memory_array_name, + ) + + assert generator.called + assert generator.call_args["owner"] is owner + assert generator.call_args["path"] == ["node1", "node2"] + assert generator.call_args["reservation"] is reservation + assert generator.call_args["memory_indices"] == [0] + assert generator.call_args["index"] == 0 + + +class NoPurificationRuleGenerator(DefaultReservationRuleGenerator): + def create_purification_rules( + self, + _owner, + _path, + _reservation, + _memory_indices, + _index, + ): + return [] + + +def test_rule_generator_can_override_only_purification_rules(): + tl = Timeline() + owner = FakeOwner("node1", tl) + generator = NoPurificationRuleGenerator() + + reservation = Reservation( + "node1", + "node2", + start_time=0, + end_time=1e12, + memory_size=1, + fidelity=0.9, + entanglement_number=1, + identity=0, + ) + + rules = generator.create_rules( + owner=owner, + path=["node1", "node2"], + reservation=reservation, + memory_indices=[0], + index=0, + ) + + actions = [rule.action for rule in rules] + + assert eg_rule_action_request in actions + assert ep_rule_action_await not in actions + assert ep_rule_action_request not in actions + assert es_rule_action_B in actions