3FS/deploy/data_placement/test/test_model.py
2025-02-27 21:53:53 +08:00

95 lines
2.7 KiB
Python

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)