feat: add default special limit

tests/special_limits.py

@@ -4,10 +4,11 @@
 """
 
 import unittest
+
 from tests.test_utils import (
     NetworkData,
-    TopologyEntry,
     NodeEntry,
+    TopologyEntry,
     execute_min_synthetic_nodes_scenario,
 )
 
@@ -131,6 +132,166 @@ def test_dfinity_limit_13_for_nns(self):
 
         assert added_providers[0] == "DFINITY"
 
+    def test_respect_default_limit(self):
+        subnet_1 = (
+            TopologyEntry()
+            .with_subnet_id("subnet_1")
+            .with_subnet_type("Application")
+            .with_country_limit(3)
+            .with_size(5)
+        )
+        subnet_2 = (
+            TopologyEntry()
+            .with_subnet_id("subnet_2")
+            .with_subnet_type("Application")
+            .with_country_limit(3)
+            .with_size(5)
+        )
+
+        aa_nodes = []
+        bb_nodes = []
+        cc_nodes = []
+
+        for _ in range(10):
+            aa_nodes.append(NodeEntry().with_country("AA"))
+            bb_nodes.append(NodeEntry().with_country("BB"))
+            cc_nodes.append(NodeEntry().with_country("CC"))
+
+        network_data = (
+            NetworkData()
+            .with_topology_entry(subnet_1)
+            .with_topology_entry(subnet_2)
+            .with_extend_nodes(aa_nodes)
+            .with_extend_nodes(bb_nodes)
+            .with_extend_nodes(cc_nodes)
+            .with_special_limit("default", "country", "AA", 2, "eq")
+            .with_special_limit("default", "country", "BB", 3, "eq")
+            .build()
+        )
+
+        output, status = execute_min_synthetic_nodes_scenario(network_data)
+
+        assert status == "Optimal"
+
+        aa_nodes_ids = [node._node_id for node in aa_nodes]
+        bb_nodes_ids = [node._node_id for node in bb_nodes]
+        cc_nodes_ids = [node._node_id for node in cc_nodes]
+
+        for subnet in output:
+            assert len(subnet["added"]) == 5
+            assert len(subnet["removed"]) == 0
+
+            added_nodes = [node["node_id"] for node in subnet["added"]]
+            # Out of added nodes only 2 are in aa
+            assert len([node for node in added_nodes if node in aa_nodes_ids]) == 2
+
+            # Out of added nodes only 3 are in bb
+            assert len([node for node in added_nodes if node in bb_nodes_ids]) == 3
+
+            # Out of added nodes 0 are in cc
+            assert len([node for node in added_nodes if node in cc_nodes_ids]) == 0
+
+    def test_append_from_default(self):
+        subnet = (
+            TopologyEntry()
+            .with_subnet_id("subnet")
+            .with_size(10)
+            .with_subnet_type("Application")
+            .with_country_limit(5)
+        )
+
+        specific_np_nodes_aa = []
+        for i in range(5):
+            specific_np_nodes_aa.append(
+                NodeEntry()
+                .with_provider_name("specific_np")
+                .with_country(f"country_{i}")
+            )
+
+        aa_country_nodes = []
+        for _ in range(5):
+            aa_country_nodes.append(NodeEntry().with_country("AA"))
+
+        other_nodes = []
+        for i in range(10):
+            other_nodes.append(NodeEntry().with_country("other_country_" + str(i)))
+
+        network_data = (
+            NetworkData()
+            .with_topology_entry(subnet)
+            .with_extend_nodes(specific_np_nodes_aa)
+            .with_extend_nodes(aa_country_nodes)
+            .with_extend_nodes(other_nodes)
+            .with_special_limit("subnet", "node_provider", "specific_np", 5, "eq")
+            .with_special_limit("default", "country", "AA", 5, "eq")
+            .build()
+        )
+
+        output, status = execute_min_synthetic_nodes_scenario(network_data)
+
+        assert status == "Optimal"
+
+        subnet_change = output[0]
+
+        added_ids = [node["node_id"] for node in subnet_change["added"]]
+
+        assert len(added_ids) == 10
+
+        specific_np_nodes_aa_ids = [node._node_id for node in specific_np_nodes_aa]
+        aa_country_nodes_ids = [node._node_id for node in aa_country_nodes]
+        # The ids used should be either aa country or specific_np
+        for id in added_ids:
+            assert id in specific_np_nodes_aa_ids or id in aa_country_nodes_ids
+
+    def test_dont_override_specific(self):
+        subnet = (
+            TopologyEntry()
+            .with_subnet_id("subnet")
+            .with_subnet_type("Application")
+            .with_size(3)
+            .with_country_limit(3)
+        )
+
+        nodes = []
+        for _ in range(2):
+            nodes.append(NodeEntry().with_country("AA"))
+            nodes.append(NodeEntry().with_country("BB"))
+
+        network_data = (
+            NetworkData()
+            .with_topology_entry(subnet)
+            .with_extend_nodes(nodes)
+            .with_special_limit("subnet", "country", "AA", 1, "eq")
+            .with_special_limit("subnet", "country", "BB", 2, "eq")
+            .with_special_limit("default", "country", "AA", 2, "eq")
+            .with_special_limit("default", "country", "BB", 1, "eq")
+            .build()
+        )
+
+        output, status = execute_min_synthetic_nodes_scenario(network_data)
+
+        assert status == "Optimal"
+
+        added_nodes = [node["node_id"] for node in output[0]["added"]]
+        assert len(added_nodes) == 3
+
+        # Expect to have 1 AA and 2 BB nodes
+        nodes_aa = [node._node_id for node in nodes if node._country == "AA"]
+        nodes_bb = [node._node_id for node in nodes if node._country == "BB"]
+
+        aa_added_nodes = 0
+        bb_added_nodes = 0
+        for node in added_nodes:
+            if node in nodes_aa:
+                aa_added_nodes += 1
+            elif node in nodes_bb:
+                bb_added_nodes += 1
+            else:
+                raise ValueError("Node is not in any seeded country")
+
+        assert aa_added_nodes == 1
+        assert bb_added_nodes == 2
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_utils.py

@@ -1,12 +1,14 @@
-from typing import List, Dict, Self, Any
 import uuid
-from topology_optimizer.utils import ALLOWED_FEATURES, parse_solver_result
+from typing import Any, Dict, List, Self
+
+import pandas as pd
+
 from topology_optimizer.data_preparation import prepare_data
 from topology_optimizer.linear_solver import (
-    solver_model_minimize_swaps,
     ATTRIBUTE_NAMES,
+    solver_model_minimize_swaps,
 )
-import pandas as pd
+from topology_optimizer.utils import ALLOWED_FEATURES, parse_solver_result
 
 
 class TopologyEntry:
@@ -125,6 +127,7 @@ class NodeEntry:
     _country: str
     _status: str
     _subnet_id: str
+    _region: str
 
     def __init__(self):
         self._node_id = str(uuid.uuid4())
@@ -164,6 +167,7 @@ def with_owner(self, owner: str) -> Self:
 
     def with_country(self, country_code: str) -> Self:
         self._region = "," + country_code + ","
+        self._country = country_code
         return self
 
     def with_status(self, status: str) -> Self:
@@ -198,7 +202,7 @@ class NetworkData:
 
     _cluster_scenario: Dict[str, List[str]]
     _cluster_scenario_name: str
-    _special_limits: Dict[int, dict[str, dict[str, (int, str)]]]
+    _special_limits: Dict[str, dict[str, dict[str, (int, str)]]]
 
     def __init__(self):
         self._cluster_scenario_name = str(uuid.uuid4())
@@ -306,25 +310,16 @@ def with_extend_sev_providers(self, providers: list[str]) -> Self:
     def with_special_limit(
         self, subnet: str, attr: str, key: str, value: int, operator: str
     ) -> Self:
-        subnet_index = None
-        for index, topology_subnet in enumerate(self._network_topology):
-            if topology_subnet._subnet_id == subnet:
-                subnet_index = index
-                break
-
-        if subnet_index is None:
-            raise ValueError(f"Subnet {subnet} not found")
-
         if attr not in ATTRIBUTE_NAMES:
             raise ValueError(f"Attribute {attr} is unknown")
 
-        if subnet_index not in self._special_limits:
-            self._special_limits[subnet_index] = {}
+        if subnet not in self._special_limits:
+            self._special_limits[subnet] = {}
 
-        if attr not in self._special_limits[subnet_index]:
-            self._special_limits[subnet_index][attr] = {}
+        if attr not in self._special_limits[subnet]:
+            self._special_limits[subnet][attr] = {}
 
-        self._special_limits[subnet_index][attr][key] = (value, operator)
+        self._special_limits[subnet][attr][key] = (value, operator)
 
         return self
 

topology_optimizer/data_preparation.py

@@ -1,13 +1,14 @@
+from typing import Any, Dict, List
+
 import pandas as pd
-from typing import Dict, Any, List
 
 from topology_optimizer.utils import (
     create_node_dataframe,
     generate_synthetic_countries,
     generate_synthetic_nodes,
     get_existing_assignment,
-    post_process_node_providers,
     mark_blacklisted_nodes,
+    post_process_node_providers,
 )
 
 
@@ -103,8 +104,10 @@ def prepare_data(
 
 def default_special_limits(
     network_topology: pd.DataFrame,
-) -> dict[int, dict[str, dict[str, (int, str)]]]:
-    nns = network_topology[network_topology["subnet_type"] == "NNS"].index[0]
+) -> dict[str, dict[str, dict[str, (int, str)]]]:
+    nns = network_topology.loc[
+        network_topology["subnet_type"] == "NNS", "subnet_id"
+    ].iloc[0]
     return {
         nns: {
             "node_provider": {"DFINITY": (3, "eq")},
@@ -113,5 +116,6 @@ def default_special_limits(
                 "Everyware": (2, "lt"),
                 "Digital Realty": (2, "lt"),
             },
-        }
+        },
+        "default": {"node_provider": {"DFINITY": (1, "eq")}},
     }

topology_optimizer/linear_solver.py

@@ -5,21 +5,23 @@
 
 """
 
+import logging
 import tempfile
+
+import pandas as pd
 from pulp import (
+    PULP_CBC_CMD,
+    LpBinary,
+    LpInteger,
+    LpMinimize,
     LpProblem,
+    LpStatus,
     LpVariable,
-    LpInteger,
-    LpBinary,
     lpSum,
-    LpMinimize,
     value,
-    LpStatus,
 )
-import pandas as pd
-from topology_optimizer.utils import get_subnet_limit
-from pulp import PULP_CBC_CMD
-import logging
+
+from topology_optimizer.utils import get_special_limits, get_subnet_limit
 
 # Standard attribute types to optimize over
 ATTRIBUTE_NAMES = [
@@ -233,10 +235,8 @@ def add_attribute_limits(network_data, model, attr):
 
     for subnet in subnet_indices:
         limit = get_subnet_limit(topology, subnet, attr)
-        special_limits = {}
-        if subnet in network_special_limits:
-            if attr in network_special_limits[subnet]:
-                special_limits = network_special_limits[subnet][attr]
+        subnet_id = topology.loc[subnet, "subnet_id"]
+        special_limits = get_special_limits(network_special_limits, subnet_id, attr)
         for idx in attr_indices:
             val = attr_list[idx]
             if val in special_limits:
@@ -253,14 +253,10 @@ def add_attribute_limits(network_data, model, attr):
                     )
                 elif op == "gt":
                     raise ValueError("`gt` doesn't make sense in our model.")
+                else:
+                    raise ValueError(f"Operation `{op}` is not supported.")
             else:
                 prob += attr_alloc[idx][subnet] <= limit, f"{attr}_Limit_{val}_{subnet}"
-                if attr == "node_provider" and val == "DFINITY":
-                    # Ensure that Dfinity is in each subnet at least once
-                    prob += (
-                        attr_alloc[idx][subnet] == 1,
-                        f"{subnet}_Subnet_Exact_1_{val}",
-                    )
 
 
 def add_attribute_subnet_allowed_values_constraints(