Initial commit

deploy/data_placement/test/test_model.py

@@ -0,0 +1,94 @@
+import copy
+import glob
+import os.path
+import importlib
+import shutil
+import tempfile
+import pytest
+from src.model.data_placement import DataPlacementModel, RebalanceTrafficModel
+
+
+placement_params = [
+  # simple cases for replication group
+  {
+    "chain_table_type": "EC",
+    "num_nodes": 5,
+    "num_targets_per_disk": 6,
+    "group_size": 2,
+  },
+  {
+    "chain_table_type": "EC",
+    "num_nodes": 5,
+    "num_targets_per_disk": 6,
+    "group_size": 3,
+  },
+  # not all targets used: num_nodes * num_targets_per_disk % group_size != 0
+  {
+    "chain_table_type": "EC",
+    "num_nodes": 7,
+    "num_targets_per_disk": 5,
+    "group_size": 4,
+  },
+  # always evenly distributed: num_targets_per_disk * (group_size-1) % (num_nodes-1) == 0
+  {
+    "chain_table_type": "EC",
+    "num_nodes": 8,
+    "num_targets_per_disk": 6,
+    "group_size": 5,
+  },
+  # all targets used & evenly distributed
+  {
+    "chain_table_type": "EC",
+    "num_nodes": 10,
+    "num_targets_per_disk": 9,
+    "group_size": 5,
+  },
+]
+qlinearize = [False, True]
+relax_lb = [1, 2]
+
+@pytest.mark.parametrize('qlinearize', qlinearize[1:])
+@pytest.mark.parametrize('relax_lb', relax_lb)
+@pytest.mark.parametrize('placement_params', placement_params)
+@pytest.mark.skipif(importlib.util.find_spec("highspy") is None, reason="cannot find solver")
+def test_solve_placement_model_with_highs(placement_params, qlinearize, relax_lb):
+  DataPlacementModel(
+    **placement_params,
+    qlinearize=qlinearize,
+    relax_lb=relax_lb,
+  ).run(pyomo_solver="appsi_highs")
+
+@pytest.mark.parametrize('chain_table_type, num_nodes, group_size', [("CR", 25, 3), ("EC", 25, 20)])
+@pytest.mark.skipif(importlib.util.find_spec("highspy") is None, reason="cannot find solver")
+def test_solve_placement_model_v25(chain_table_type, num_nodes, group_size):
+  model = DataPlacementModel(
+    chain_table_type=chain_table_type,
+    num_nodes=num_nodes,
+    group_size=group_size,
+    qlinearize=True,
+    relax_lb=1,
+    relax_ub=1,
+  )
+  model.run(pyomo_solver="appsi_highs", max_timelimit=30, auto_relax=True)
+
+@pytest.mark.parametrize('placement_params', placement_params)
+@pytest.mark.skipif(importlib.util.find_spec("highspy") is None, reason="cannot find solver")
+def test_solve_rebalance_model(placement_params):
+  model = DataPlacementModel(
+    **placement_params,
+    qlinearize=True,
+    relax_lb=1,
+    relax_ub=1,
+  )
+  instance = model.run(pyomo_solver="appsi_highs")
+
+  placement_params = copy.deepcopy(placement_params)
+  placement_params["num_nodes"] *= 2
+  placement_params.pop("num_targets_per_disk")
+  RebalanceTrafficModel(
+    existing_incidence_matrix=model.get_incidence_matrix(instance),
+    **placement_params,
+    qlinearize=True,
+    relax_lb=2,
+    relax_ub=1,
+  ).run(pyomo_solver="appsi_highs", max_timelimit=15, auto_relax=True)

deploy/data_placement/test/test_plan.py

@@ -0,0 +1,10 @@
+from smallpond.test_fabric import TestFabric
+from src.model.data_placement_job import search_data_placement_plans
+
+class TestPlan(TestFabric):
+
+  def test_search_data_placement_plans(self):
+    for pyomo_solver in ["appsi_highs"]:
+      with self.subTest(pyomo_solver=pyomo_solver):
+        plan = search_data_placement_plans(chain_table_type="EC", num_nodes=[10], group_size=[5, 9], solver_threads=16, pyomo_solver=pyomo_solver)
+        self.execute_plan(plan, num_executors=1)

deploy/data_placement/test/test_setup.py

@@ -0,0 +1,55 @@
+from collections import Counter
+import glob
+import os.path
+import pytest
+
+from src.model.data_placement import DataPlacementModel
+from src.setup.gen_chain_table import generate_chains
+
+
+@pytest.mark.parametrize('num_nodes, num_disks_per_node, num_targets_per_disk, num_replicas', [(5, 10, 6, 2), (10, 10, 9, 3)])
+def test_generate_cr_chains(num_nodes: int, num_disks_per_node: int, num_targets_per_disk: int, num_replicas: int):
+  model = DataPlacementModel(
+    chain_table_type="CR",
+    num_nodes=num_nodes,
+    num_targets_per_disk=num_targets_per_disk,
+    group_size=num_replicas,
+    qlinearize=True,
+    relax_lb=1,
+    relax_ub=1,
+  )
+  instance = model.run(pyomo_solver="appsi_highs", max_timelimit=15, auto_relax=True)
+
+  generate_chains(
+    chain_table_type="CR",
+    node_id_begin=1,
+    node_id_end=num_nodes,
+    num_disks_per_node=num_disks_per_node,
+    num_targets_per_disk=num_targets_per_disk,
+    target_id_prefix=1,
+    chain_id_prefix=9,
+    incidence_matrix=model.get_incidence_matrix(instance))
+
+
+@pytest.mark.parametrize('num_nodes, num_disks_per_node, num_targets_per_disk, ec_group_size', [(20, 10, 6, 12), (25, 10, 12, 20)])
+def test_generate_ec_chains(num_nodes: int, num_disks_per_node: int, num_targets_per_disk: int, ec_group_size: int):
+  model = DataPlacementModel(
+    chain_table_type="EC",
+    num_nodes=num_nodes,
+    num_targets_per_disk=num_targets_per_disk,
+    group_size=ec_group_size,
+    qlinearize=True,
+    relax_lb=1,
+    relax_ub=1,
+  )
+  instance = model.run(pyomo_solver="appsi_highs", max_timelimit=15, auto_relax=True)
+
+  generate_chains(
+    chain_table_type="EC",
+    node_id_begin=1,
+    node_id_end=num_nodes,
+    num_disks_per_node=num_disks_per_node,
+    num_targets_per_disk=num_targets_per_disk,
+    target_id_prefix=1,
+    chain_id_prefix=9,
+    incidence_matrix=model.get_incidence_matrix(instance))