From 8a20fc679b02a5276819550ab3d21372763687dd Mon Sep 17 00:00:00 2001 From: tvk134 Date: Sat, 23 May 2026 17:04:04 -0500 Subject: [PATCH 1/3] Add registry for reservation rule classes --- .../reservation_rule_registry.py | 286 ++++++++++++++++++ .../resource_management/resource_manager.py | 108 +------ .../test_reservation_rule_registry.py | 76 +++++ 3 files changed, 376 insertions(+), 94 deletions(-) create mode 100644 sequence/resource_management/reservation_rule_registry.py create mode 100644 tests/resource_management/test_reservation_rule_registry.py diff --git a/sequence/resource_management/reservation_rule_registry.py b/sequence/resource_management/reservation_rule_registry.py new file mode 100644 index 000000000..b0bae9509 --- /dev/null +++ b/sequence/resource_management/reservation_rule_registry.py @@ -0,0 +1,286 @@ +"""Registry utilities for reservation-generated resource-management rules. + +This module defines reservation-specific Rule subclasses and a registry that +maps reservation rule names to Rule constructors. The default resource-management +rule generator uses this registry to construct the built-in rules, while users +may replace or disable individual rules without subclassing the generator. +""" + +from typing import Final + +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 + + +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, +) + +ReservationRuleConstructor = type[Rule] + + +class EGAwaitRule(Rule): + """Entanglement-generation await rule for non-left-end nodes.""" + + def __init__(self, owner, path, reservation, memory_indices, index) -> None: + condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} + action_args = { + "mid": owner.map_to_middle_node[path[index - 1]], + "path": path, + "index": index, + } + super().__init__(10, eg_rule_action_await, eg_rule_condition, action_args, condition_args) + + +class EGRequestRule(Rule): + """Entanglement-generation request rule for non-right-end nodes.""" + + def __init__(self, owner, path, reservation, memory_indices, index) -> None: + 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, + } + super().__init__(10, eg_rule_action_request, eg_rule_condition, action_args, condition_args) + + +class EPRequestRule(Rule): + """Entanglement-purification request rule.""" + + def __init__(self, _owner, _path, reservation, memory_indices, _index) -> None: + condition_args = { + "memory_indices": memory_indices[:reservation.memory_size], + "reservation": reservation, + "purification_mode": reservation.purification_mode, + } + super().__init__(10, ep_rule_action_request, ep_rule_condition_request, {}, condition_args) + + +class EPAwaitRule(Rule): + """Entanglement-purification await rule.""" + + def __init__(self, _owner, path, reservation, memory_indices, index) -> None: + 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, + } + + super().__init__(10, ep_rule_action_await, ep_rule_condition_await, {}, condition_args) + + +class ESBEndRule(Rule): + """Endpoint entanglement-swapping B rule.""" + + def __init__(self, _owner, path, reservation, memory_indices, index) -> None: + if index == 0: + target_remote = path[-1] + else: + target_remote = path[0] + + condition_args = { + "memory_indices": memory_indices, + "target_remote": target_remote, + "fidelity": reservation.fidelity, + } + super().__init__(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] + + +class ESARule(Rule): + """Middle-node entanglement-swapping A rule.""" + + def __init__(self, owner, path, reservation, memory_indices, _index) -> None: + left, right = _get_swapping_neighbors(owner.name, path) + condition_args = { + "memory_indices": memory_indices, + "left": left, + "right": right, + "fidelity": reservation.fidelity, + } + super().__init__(10, es_rule_action_A, es_rule_condition_A, {}, condition_args) + + +class ESBRule(Rule): + """Middle-node entanglement-swapping B rule.""" + + def __init__(self, owner, path, reservation, memory_indices, _index) -> None: + left, right = _get_swapping_neighbors(owner.name, path) + condition_args = { + "memory_indices": memory_indices, + "left": left, + "right": right, + "fidelity": reservation.fidelity, + } + super().__init__(10, es_rule_action_B, es_rule_condition_B, {}, condition_args) + + +DEFAULT_RESERVATION_RULE_CONSTRUCTORS: Final[dict[str, ReservationRuleConstructor]] = { + EG_AWAIT: EGAwaitRule, + EG_REQUEST: EGRequestRule, + EP_REQUEST: EPRequestRule, + EP_AWAIT: EPAwaitRule, + ES_B_END: ESBEndRule, + ES_A: ESARule, + ES_B: ESBRule, +} + +DEFAULT_RESERVATION_RULE_ORDER: Final[tuple[str, ...]] = ( + EG_AWAIT, + EG_REQUEST, + EP_REQUEST, + EP_AWAIT, + ES_B_END, + ES_A, + ES_B, +) + + +class ReservationRuleRegistry: + """Registry for reservation Rule constructors.""" + + def __init__(self, constructors: dict[str, ReservationRuleConstructor] | None = None) -> None: + self._constructors: dict[str, ReservationRuleConstructor] = {} + if constructors is not None: + for name, constructor in constructors.items(): + self.register(name, constructor) + + def register(self, name: str, constructor: ReservationRuleConstructor) -> None: + """Register or replace a reservation Rule constructor.""" + self._validate_name(name) + if not issubclass(constructor, Rule): + msg = "Reservation rule constructors must be Rule subclasses" + raise TypeError(msg) + self._constructors[name] = constructor + + def disable(self, name: str) -> None: + """Disable a reservation Rule constructor.""" + self._validate_name(name) + self._constructors.pop(name, None) + + def get(self, name: str) -> ReservationRuleConstructor | None: + """Return the constructor registered for a rule name, if any.""" + self._validate_name(name) + return self._constructors.get(name) + + def build(self, name: str, *args, **kwargs) -> Rule | None: + """Build a rule by name. + + Returns None when the rule is disabled. + """ + constructor = self.get(name) + if constructor is None: + return None + return constructor(*args, **kwargs) + + def copy(self) -> "ReservationRuleRegistry": + """Return a shallow copy of this registry.""" + return ReservationRuleRegistry(self._constructors.copy()) + + @staticmethod + def _validate_name(name: str) -> None: + if name not in RESERVATION_RULE_NAMES: + msg = f"Unknown reservation rule name: {name}" + raise ValueError(msg) + + +class DefaultReservationRuleGenerator: + """Default generator for reservation-based resource-management rules.""" + + def __init__(self, registry: ReservationRuleRegistry | None = None) -> None: + if registry is None: + registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) + self.registry = registry + + def create_rules(self, owner, path, reservation, memory_indices, index) -> list[Rule]: + """Create reservation rules using the configured rule registry.""" + rules: list[Rule] = [] + + if index > 0: + rule = self.registry.build(EG_AWAIT, owner, path, reservation, memory_indices, index) + if rule is not None: + rules.append(rule) + + if index < len(path) - 1: + rule = self.registry.build(EG_REQUEST, owner, path, reservation, memory_indices, index) + if rule is not None: + rules.append(rule) + + if index > 0: + rule = self.registry.build(EP_REQUEST, owner, path, reservation, memory_indices, index) + if rule is not None: + rules.append(rule) + + if index < len(path) - 1: + rule = self.registry.build(EP_AWAIT, owner, path, reservation, memory_indices, index) + if rule is not None: + rules.append(rule) + + if index == 0 or index == len(path) - 1: + rule = self.registry.build(ES_B_END, owner, path, reservation, memory_indices, index) + if rule is not None: + rules.append(rule) + else: + rule = self.registry.build(ES_A, owner, path, reservation, memory_indices, index) + if rule is not None: + rules.append(rule) + + rule = self.registry.build(ES_B, owner, path, reservation, memory_indices, index) + 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..4e9992639 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 ( + DefaultReservationRuleGenerator, + ReservationRuleRegistry, ) from ..kernel.event import Event from ..kernel.process import Process @@ -148,9 +138,11 @@ 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 = DefaultReservationRuleGenerator() - def init(self): - pass + def get_reservation_rule_registry(self) -> ReservationRuleRegistry: + """Return the registry used to build reservation rules.""" + 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. @@ -169,85 +161,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..621502ffd --- /dev/null +++ b/tests/resource_management/test_reservation_rule_registry.py @@ -0,0 +1,76 @@ +from sequence.resource_management.reservation_rule_registry import ( + DEFAULT_RESERVATION_RULE_CONSTRUCTORS, + EGRequestRule, + EG_REQUEST, + EPAwaitRule, + EP_AWAIT, + EPRequestRule, + EP_REQUEST, + DefaultReservationRuleGenerator, + ReservationRuleRegistry, +) +from sequence.resource_management.rule_manager import Rule + + +def test_registry_can_disable_individual_rule_constructor(): + registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) + + assert registry.get(EP_REQUEST) is EPRequestRule + + registry.disable(EP_REQUEST) + + assert registry.get(EP_REQUEST) is None + assert registry.get(EG_REQUEST) is EGRequestRule + assert registry.get(EP_AWAIT) is EPAwaitRule + + +def test_registry_can_replace_individual_rule_constructor(): + registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) + + class CustomEPRequestRule(Rule): + def __init__(self, *_args, **_kwargs) -> None: + super().__init__(10, lambda *_: [], lambda *_: [], {}, {}) + + registry.register(EP_REQUEST, CustomEPRequestRule) + + assert registry.get(EP_REQUEST) is CustomEPRequestRule + assert registry.get(EG_REQUEST) is EGRequestRule + + +def test_default_generator_uses_registry_to_disable_rule(): + registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) + generator = DefaultReservationRuleGenerator(registry) + + class Owner: + name = "node1" + map_to_middle_node = {"node2": "middle"} + + class Reservation: + memory_size = 1 + fidelity = 0.9 + purification_mode = "BBPSSW" + + rules = generator.create_rules( + Owner(), + ["node1", "node2"], + Reservation(), + [0], + 0, + ) + + assert all(isinstance(rule, Rule) for rule in rules) + assert any(isinstance(rule, EGRequestRule) for rule in rules) + assert any(isinstance(rule, EPAwaitRule) for rule in rules) + + registry.disable(EP_AWAIT) + + rules = generator.create_rules( + Owner(), + ["node1", "node2"], + Reservation(), + [0], + 0, + ) + + assert any(isinstance(rule, EGRequestRule) for rule in rules) + assert not any(isinstance(rule, EPAwaitRule) for rule in rules) From 15d0538d9712b5ec485ef9fd10a74db3af33b76d Mon Sep 17 00:00:00 2001 From: tvk134 Date: Tue, 2 Jun 2026 20:53:08 -0500 Subject: [PATCH 2/3] Refine reservation rule registry API --- .../reservation_rule_registry.py | 94 +++++++++++++++---- .../resource_management/resource_manager.py | 4 +- .../test_reservation_rule_registry.py | 50 +++++++--- 3 files changed, 112 insertions(+), 36 deletions(-) diff --git a/sequence/resource_management/reservation_rule_registry.py b/sequence/resource_management/reservation_rule_registry.py index b0bae9509..6b9eb2c00 100644 --- a/sequence/resource_management/reservation_rule_registry.py +++ b/sequence/resource_management/reservation_rule_registry.py @@ -49,20 +49,36 @@ class EGAwaitRule(Rule): """Entanglement-generation await rule for non-left-end nodes.""" - def __init__(self, owner, path, reservation, memory_indices, index) -> None: + def __init__( + self, + owner, + path, + reservation, + memory_indices, + index, + priority: int = 10, + ) -> None: condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} action_args = { "mid": owner.map_to_middle_node[path[index - 1]], "path": path, "index": index, } - super().__init__(10, eg_rule_action_await, eg_rule_condition, action_args, condition_args) + super().__init__(priority, eg_rule_action_await, eg_rule_condition, action_args, condition_args) class EGRequestRule(Rule): """Entanglement-generation request rule for non-right-end nodes.""" - def __init__(self, owner, path, reservation, memory_indices, index) -> None: + def __init__( + self, + owner, + path, + reservation, + memory_indices, + index, + priority: int = 10, + ) -> None: if index == 0: condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} else: @@ -75,25 +91,41 @@ def __init__(self, owner, path, reservation, memory_indices, index) -> None: "name": owner.name, "reservation": reservation, } - super().__init__(10, eg_rule_action_request, eg_rule_condition, action_args, condition_args) + super().__init__(priority, eg_rule_action_request, eg_rule_condition, action_args, condition_args) class EPRequestRule(Rule): """Entanglement-purification request rule.""" - def __init__(self, _owner, _path, reservation, memory_indices, _index) -> None: + def __init__( + self, + _owner, + _path, + reservation, + memory_indices, + _index, + priority: int = 10, + ) -> None: condition_args = { "memory_indices": memory_indices[:reservation.memory_size], "reservation": reservation, "purification_mode": reservation.purification_mode, } - super().__init__(10, ep_rule_action_request, ep_rule_condition_request, {}, condition_args) + super().__init__(priority, ep_rule_action_request, ep_rule_condition_request, {}, condition_args) class EPAwaitRule(Rule): """Entanglement-purification await rule.""" - def __init__(self, _owner, path, reservation, memory_indices, index) -> None: + def __init__( + self, + _owner, + path, + reservation, + memory_indices, + index, + priority: int = 10, + ) -> None: if index == 0: condition_args = { "memory_indices": memory_indices, @@ -107,13 +139,21 @@ def __init__(self, _owner, path, reservation, memory_indices, index) -> None: "purification_mode": reservation.purification_mode, } - super().__init__(10, ep_rule_action_await, ep_rule_condition_await, {}, condition_args) + super().__init__(priority, ep_rule_action_await, ep_rule_condition_await, {}, condition_args) class ESBEndRule(Rule): """Endpoint entanglement-swapping B rule.""" - def __init__(self, _owner, path, reservation, memory_indices, index) -> None: + def __init__( + self, + _owner, + path, + reservation, + memory_indices, + index, + priority: int = 10, + ) -> None: if index == 0: target_remote = path[-1] else: @@ -124,7 +164,7 @@ def __init__(self, _owner, path, reservation, memory_indices, index) -> None: "target_remote": target_remote, "fidelity": reservation.fidelity, } - super().__init__(10, es_rule_action_B, es_rule_condition_B_end, {}, condition_args) + super().__init__(priority, es_rule_action_B, es_rule_condition_B_end, {}, condition_args) def _get_swapping_neighbors(owner_name: str, path: list[str]) -> tuple[str, str]: @@ -144,7 +184,15 @@ def _get_swapping_neighbors(owner_name: str, path: list[str]) -> tuple[str, str] class ESARule(Rule): """Middle-node entanglement-swapping A rule.""" - def __init__(self, owner, path, reservation, memory_indices, _index) -> None: + def __init__( + self, + owner, + path, + reservation, + memory_indices, + _index, + priority: int = 10, + ) -> None: left, right = _get_swapping_neighbors(owner.name, path) condition_args = { "memory_indices": memory_indices, @@ -152,13 +200,21 @@ def __init__(self, owner, path, reservation, memory_indices, _index) -> None: "right": right, "fidelity": reservation.fidelity, } - super().__init__(10, es_rule_action_A, es_rule_condition_A, {}, condition_args) + super().__init__(priority, es_rule_action_A, es_rule_condition_A, {}, condition_args) class ESBRule(Rule): """Middle-node entanglement-swapping B rule.""" - def __init__(self, owner, path, reservation, memory_indices, _index) -> None: + def __init__( + self, + owner, + path, + reservation, + memory_indices, + _index, + priority: int = 10, + ) -> None: left, right = _get_swapping_neighbors(owner.name, path) condition_args = { "memory_indices": memory_indices, @@ -166,7 +222,7 @@ def __init__(self, owner, path, reservation, memory_indices, _index) -> None: "right": right, "fidelity": reservation.fidelity, } - super().__init__(10, es_rule_action_B, es_rule_condition_B, {}, condition_args) + super().__init__(priority, es_rule_action_B, es_rule_condition_B, {}, condition_args) DEFAULT_RESERVATION_RULE_CONSTRUCTORS: Final[dict[str, ReservationRuleConstructor]] = { @@ -238,13 +294,11 @@ def _validate_name(name: str) -> None: raise ValueError(msg) -class DefaultReservationRuleGenerator: - """Default generator for reservation-based resource-management rules.""" +class ReservationRuleGenerator: + """Rule generator for reservation-based resource-management rules.""" - def __init__(self, registry: ReservationRuleRegistry | None = None) -> None: - if registry is None: - registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) - self.registry = registry + def __init__(self) -> None: + self.registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) def create_rules(self, owner, path, reservation, memory_indices, index) -> list[Rule]: """Create reservation rules using the configured rule registry.""" diff --git a/sequence/resource_management/resource_manager.py b/sequence/resource_management/resource_manager.py index 4e9992639..84c8e4202 100644 --- a/sequence/resource_management/resource_manager.py +++ b/sequence/resource_management/resource_manager.py @@ -11,7 +11,7 @@ from collections.abc import Callable from .reservation_rule_registry import ( - DefaultReservationRuleGenerator, + ReservationRuleGenerator, ReservationRuleRegistry, ) from ..kernel.event import Event @@ -138,7 +138,7 @@ 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 = DefaultReservationRuleGenerator() + self.rule_generator = ReservationRuleGenerator() def get_reservation_rule_registry(self) -> ReservationRuleRegistry: """Return the registry used to build reservation rules.""" diff --git a/tests/resource_management/test_reservation_rule_registry.py b/tests/resource_management/test_reservation_rule_registry.py index 621502ffd..eb7b2abe9 100644 --- a/tests/resource_management/test_reservation_rule_registry.py +++ b/tests/resource_management/test_reservation_rule_registry.py @@ -6,12 +6,23 @@ EP_AWAIT, EPRequestRule, EP_REQUEST, - DefaultReservationRuleGenerator, + ReservationRuleGenerator, ReservationRuleRegistry, ) from sequence.resource_management.rule_manager import Rule +class Owner: + name = "node1" + map_to_middle_node = {"node2": "middle"} + + +class Reservation: + memory_size = 1 + fidelity = 0.9 + purification_mode = "BBPSSW" + + def test_registry_can_disable_individual_rule_constructor(): registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) @@ -27,9 +38,8 @@ def test_registry_can_disable_individual_rule_constructor(): def test_registry_can_replace_individual_rule_constructor(): registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) - class CustomEPRequestRule(Rule): - def __init__(self, *_args, **_kwargs) -> None: - super().__init__(10, lambda *_: [], lambda *_: [], {}, {}) + class CustomEPRequestRule(EPRequestRule): + pass registry.register(EP_REQUEST, CustomEPRequestRule) @@ -37,18 +47,30 @@ def __init__(self, *_args, **_kwargs) -> None: assert registry.get(EG_REQUEST) is EGRequestRule -def test_default_generator_uses_registry_to_disable_rule(): - registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) - generator = DefaultReservationRuleGenerator(registry) +def test_generator_uses_registered_replacement_rule(): + generator = ReservationRuleGenerator() + + class CustomEGRequestRule(EGRequestRule): + pass + + generator.registry.register(EG_REQUEST, CustomEGRequestRule) + + rules = generator.create_rules( + Owner(), + ["node1", "node2"], + Reservation(), + [0], + 0, + ) + + assert all(isinstance(rule, Rule) for rule in rules) + assert any(isinstance(rule, CustomEGRequestRule) for rule in rules) + assert not any(type(rule) is EGRequestRule for rule in rules) - class Owner: - name = "node1" - map_to_middle_node = {"node2": "middle"} - class Reservation: - memory_size = 1 - fidelity = 0.9 - purification_mode = "BBPSSW" +def test_generator_uses_registry_to_disable_rule(): + generator = ReservationRuleGenerator() + registry = generator.registry rules = generator.create_rules( Owner(), From ef10699714d634b0b39eaba4f0b53420925ed38f Mon Sep 17 00:00:00 2001 From: tvk134 Date: Sat, 11 Jul 2026 22:22:18 +0800 Subject: [PATCH 3/3] Simplify reservation rule registry scope --- .../reservation_rule_registry.py | 453 ++++++++---------- .../resource_management/resource_manager.py | 6 +- .../test_reservation_rule_registry.py | 129 +++-- 3 files changed, 313 insertions(+), 275 deletions(-) diff --git a/sequence/resource_management/reservation_rule_registry.py b/sequence/resource_management/reservation_rule_registry.py index 6b9eb2c00..c30de0bf6 100644 --- a/sequence/resource_management/reservation_rule_registry.py +++ b/sequence/resource_management/reservation_rule_registry.py @@ -1,12 +1,17 @@ -"""Registry utilities for reservation-generated resource-management rules. +"""Utilities for building reservation-generated resource-management rules. -This module defines reservation-specific Rule subclasses and a registry that -maps reservation rule names to Rule constructors. The default resource-management -rule generator uses this registry to construct the built-in rules, while users -may replace or disable individual rules without subclassing the generator. +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 typing import Final +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, @@ -24,6 +29,10 @@ ) 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" @@ -43,128 +52,119 @@ ES_B, ) -ReservationRuleConstructor = type[Rule] - - -class EGAwaitRule(Rule): - """Entanglement-generation await rule for non-left-end nodes.""" - - def __init__( - self, - owner, - path, - reservation, - memory_indices, - index, - priority: int = 10, - ) -> None: - condition_args = {"memory_indices": memory_indices[:reservation.memory_size]} - action_args = { - "mid": owner.map_to_middle_node[path[index - 1]], - "path": path, - "index": index, - } - super().__init__(priority, eg_rule_action_await, eg_rule_condition, action_args, condition_args) - - -class EGRequestRule(Rule): - """Entanglement-generation request rule for non-right-end nodes.""" - - def __init__( - self, - owner, - path, - reservation, - memory_indices, - index, - priority: int = 10, - ) -> None: - 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, - } - super().__init__(priority, eg_rule_action_request, eg_rule_condition, action_args, condition_args) +@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) -class EPRequestRule(Rule): - """Entanglement-purification request rule.""" - def __init__( - self, - _owner, - _path, - reservation, - memory_indices, - _index, - priority: int = 10, - ) -> None: +def _build_ep_await_rule(context: ReservationRuleContext) -> Rule: + if context.index == 0: condition_args = { - "memory_indices": memory_indices[:reservation.memory_size], - "reservation": reservation, - "purification_mode": reservation.purification_mode, + "memory_indices": context.memory_indices, + "fidelity": context.reservation.fidelity, + "purification_mode": context.reservation.purification_mode, } - super().__init__(priority, ep_rule_action_request, ep_rule_condition_request, {}, condition_args) - - -class EPAwaitRule(Rule): - """Entanglement-purification await rule.""" - - def __init__( - self, - _owner, - path, - reservation, - memory_indices, - index, - priority: int = 10, - ) -> None: - 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, - } - - super().__init__(priority, ep_rule_action_await, ep_rule_condition_await, {}, condition_args) - - -class ESBEndRule(Rule): - """Endpoint entanglement-swapping B rule.""" - - def __init__( - self, - _owner, - path, - reservation, - memory_indices, - index, - priority: int = 10, - ) -> None: - if index == 0: - target_remote = path[-1] - else: - target_remote = path[0] - + else: condition_args = { - "memory_indices": memory_indices, - "target_remote": target_remote, - "fidelity": reservation.fidelity, + "memory_indices": context.memory_indices[context.reservation.memory_size:], + "fidelity": context.reservation.fidelity, + "purification_mode": context.reservation.purification_mode, } - super().__init__(priority, es_rule_action_B, es_rule_condition_B_end, {}, condition_args) + + 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]: @@ -181,111 +181,88 @@ def _get_swapping_neighbors(owner_name: str, path: list[str]) -> tuple[str, str] return reduced_path[reduced_index - 1], reduced_path[reduced_index + 1] -class ESARule(Rule): - """Middle-node entanglement-swapping A rule.""" - - def __init__( - self, - owner, - path, - reservation, - memory_indices, - _index, - priority: int = 10, - ) -> None: - left, right = _get_swapping_neighbors(owner.name, path) - condition_args = { - "memory_indices": memory_indices, - "left": left, - "right": right, - "fidelity": reservation.fidelity, - } - super().__init__(priority, es_rule_action_A, es_rule_condition_A, {}, condition_args) - - -class ESBRule(Rule): - """Middle-node entanglement-swapping B rule.""" - - def __init__( - self, - owner, - path, - reservation, - memory_indices, - _index, - priority: int = 10, - ) -> None: - left, right = _get_swapping_neighbors(owner.name, path) - condition_args = { - "memory_indices": memory_indices, - "left": left, - "right": right, - "fidelity": reservation.fidelity, - } - super().__init__(priority, es_rule_action_B, es_rule_condition_B, {}, condition_args) - - -DEFAULT_RESERVATION_RULE_CONSTRUCTORS: Final[dict[str, ReservationRuleConstructor]] = { - EG_AWAIT: EGAwaitRule, - EG_REQUEST: EGRequestRule, - EP_REQUEST: EPRequestRule, - EP_AWAIT: EPAwaitRule, - ES_B_END: ESBEndRule, - ES_A: ESARule, - ES_B: ESBRule, -} - -DEFAULT_RESERVATION_RULE_ORDER: Final[tuple[str, ...]] = ( - EG_AWAIT, - EG_REQUEST, - EP_REQUEST, - EP_AWAIT, - ES_B_END, - ES_A, - ES_B, +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 reservation Rule constructors.""" - - def __init__(self, constructors: dict[str, ReservationRuleConstructor] | None = None) -> None: - self._constructors: dict[str, ReservationRuleConstructor] = {} - if constructors is not None: - for name, constructor in constructors.items(): - self.register(name, constructor) - - def register(self, name: str, constructor: ReservationRuleConstructor) -> None: - """Register or replace a reservation Rule constructor.""" + """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) - if not issubclass(constructor, Rule): - msg = "Reservation rule constructors must be Rule subclasses" - raise TypeError(msg) - self._constructors[name] = constructor + self._validate_builder(builder) + self._builders[name] = builder def disable(self, name: str) -> None: - """Disable a reservation Rule constructor.""" + """Disable an existing reservation rule slot.""" self._validate_name(name) - self._constructors.pop(name, None) + self._builders.pop(name, None) - def get(self, name: str) -> ReservationRuleConstructor | None: - """Return the constructor registered for a rule name, if any.""" + 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._constructors.get(name) + return self._builders.get(name) - def build(self, name: str, *args, **kwargs) -> Rule | None: - """Build a rule by name. + def build(self, name: str, context: ReservationRuleContext) -> Rule | None: + """Build the named rule for the supplied context. - Returns None when the rule is disabled. + Returns None when the rule slot is disabled. """ - constructor = self.get(name) - if constructor is None: + builder = self.get(name) + if builder is None: return None - return constructor(*args, **kwargs) - def copy(self) -> "ReservationRuleRegistry": - """Return a shallow copy of this registry.""" - return ReservationRuleRegistry(self._constructors.copy()) + 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: @@ -293,48 +270,34 @@ def _validate_name(name: str) -> None: 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(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) + self.registry = ReservationRuleRegistry() def create_rules(self, owner, path, reservation, memory_indices, index) -> list[Rule]: - """Create reservation rules using the configured rule registry.""" + """Create reservation rules in SeQUeNCe's static rule-slot order.""" + context = ReservationRuleContext(owner, path, reservation, memory_indices, index) rules: list[Rule] = [] - if index > 0: - rule = self.registry.build(EG_AWAIT, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(rule) - - if index < len(path) - 1: - rule = self.registry.build(EG_REQUEST, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(rule) - - if index > 0: - rule = self.registry.build(EP_REQUEST, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(rule) - - if index < len(path) - 1: - rule = self.registry.build(EP_AWAIT, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(rule) - - if index == 0 or index == len(path) - 1: - rule = self.registry.build(ES_B_END, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(rule) - else: - rule = self.registry.build(ES_A, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(rule) - - rule = self.registry.build(ES_B, owner, path, reservation, memory_indices, index) - if rule is not None: - rules.append(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 84c8e4202..ca6d9707a 100644 --- a/sequence/resource_management/resource_manager.py +++ b/sequence/resource_management/resource_manager.py @@ -140,8 +140,11 @@ def __init__(self, owner: QuantumRouter, memory_array_name: str): 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 build reservation rules.""" + """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): @@ -153,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: diff --git a/tests/resource_management/test_reservation_rule_registry.py b/tests/resource_management/test_reservation_rule_registry.py index eb7b2abe9..6c90cedbf 100644 --- a/tests/resource_management/test_reservation_rule_registry.py +++ b/tests/resource_management/test_reservation_rule_registry.py @@ -1,11 +1,20 @@ +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_CONSTRUCTORS, - EGRequestRule, + DEFAULT_RESERVATION_RULE_SPECS, + EG_AWAIT, EG_REQUEST, - EPAwaitRule, EP_AWAIT, - EPRequestRule, EP_REQUEST, + ES_A, + ES_B, + ES_B_END, + ReservationRuleContext, ReservationRuleGenerator, ReservationRuleRegistry, ) @@ -14,7 +23,9 @@ class Owner: name = "node1" - map_to_middle_node = {"node2": "middle"} + map_to_middle_node = {"node0": "middle_left", "node2": "middle_right"} + swapping_success_prob = 0.75 + swapping_degradation = 0.95 class Reservation: @@ -23,37 +34,51 @@ class Reservation: purification_mode = "BBPSSW" -def test_registry_can_disable_individual_rule_constructor(): - registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) +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 EPRequestRule + 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 EGRequestRule - assert registry.get(EP_AWAIT) is EPAwaitRule + assert registry.get(EG_REQUEST) is not None + assert registry.get(EP_AWAIT) is not None -def test_registry_can_replace_individual_rule_constructor(): - registry = ReservationRuleRegistry(DEFAULT_RESERVATION_RULE_CONSTRUCTORS) +def test_registry_can_replace_individual_rule_builder(): + registry = ReservationRuleRegistry() - class CustomEPRequestRule(EPRequestRule): - pass + registry.replace(EP_REQUEST, _custom_rule_builder) - registry.register(EP_REQUEST, CustomEPRequestRule) + assert registry.get(EP_REQUEST) is _custom_rule_builder + assert registry.get(EG_REQUEST) is not _custom_rule_builder - assert registry.get(EP_REQUEST) is CustomEPRequestRule - assert registry.get(EG_REQUEST) is EGRequestRule +def test_registry_rejects_unknown_rule_name(): + registry = ReservationRuleRegistry() -def test_generator_uses_registered_replacement_rule(): - generator = ReservationRuleGenerator() + 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: [], {}, {}) - class CustomEGRequestRule(EGRequestRule): - pass + with pytest.raises(TypeError): + registry.replace(EP_REQUEST, bad_builder) - generator.registry.register(EG_REQUEST, CustomEGRequestRule) + +def test_generator_uses_replacement_builder(): + generator = ReservationRuleGenerator() + generator.registry.replace(EG_REQUEST, _custom_rule_builder) rules = generator.create_rules( Owner(), @@ -64,8 +89,8 @@ class CustomEGRequestRule(EGRequestRule): ) assert all(isinstance(rule, Rule) for rule in rules) - assert any(isinstance(rule, CustomEGRequestRule) for rule in rules) - assert not any(type(rule) is EGRequestRule 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(): @@ -81,8 +106,8 @@ def test_generator_uses_registry_to_disable_rule(): ) assert all(isinstance(rule, Rule) for rule in rules) - assert any(isinstance(rule, EGRequestRule) for rule in rules) - assert any(isinstance(rule, EPAwaitRule) 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) @@ -94,5 +119,53 @@ def test_generator_uses_registry_to_disable_rule(): 0, ) - assert any(isinstance(rule, EGRequestRule) for rule in rules) - assert not any(isinstance(rule, EPAwaitRule) for rule in rules) + 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"