init v0.11.0rc0

tests/ut/eplb/core/policy/test_policy_abstract.py

@@ -0,0 +1,31 @@
+# test_policy_abstract.py
+from vllm_ascend.eplb.core.policy.policy_abstract import (DynamicConfig,
+                                                          EplbPolicy)
+
+
+class DummyPolicy(EplbPolicy):
+
+    def rebalance_experts(self, current_expert_table, expert_workload):
+        return 1, current_expert_table
+
+
+def test_dynamic_config_attributes():
+    config = DynamicConfig()
+    assert config.placement_policy is None
+    assert config.max_transferred_expert_per_layer == 100
+    assert config.ep_worldsize == 64
+    assert config.num_die_per_host == 8
+
+
+def test_eplb_policy_init_and_method():
+    config = DynamicConfig()
+    policy = DummyPolicy(config)
+
+    assert policy.config == config
+
+    expert_table = [[0, 1, 2]]
+    workload = [10]
+    res, new_table = policy.rebalance_experts(expert_table, workload)
+
+    assert res == 1
+    assert new_table == expert_table

tests/ut/eplb/core/policy/test_policy_dynamic_ep.py

@@ -0,0 +1,98 @@
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from vllm_ascend.eplb.core.policy.policy_dynamic_ep import DynamicEplb
+
+
+class TestDynamicEplb:
+
+    def test_add_redundant_basic(self):
+        current_expert_table = np.array([[[0, 1], [1, 0]]])
+        expert_workload = np.array([[[2, 3], [4, 1]]])
+        num_original_expert = 2
+        result = DynamicEplb.add_redundant(current_expert_table,
+                                           expert_workload,
+                                           num_original_expert)
+        expected = np.array([[2 + 1, 3 + 4]])
+        assert np.array_equal(result, expected)
+
+    def test_get_redundant_num(self):
+        counts = np.array([2, 1, 3])
+        assert DynamicEplb.get_redundant_num(3, counts) == 3
+
+    def test_calculate_max_heat_per_layer(self):
+        workload_table = np.array([[[1, 2], [3, 4]], [[2, 2], [1, 1]]])
+        max_heat = DynamicEplb.calculate_max_heat_per_layer(workload_table, 2)
+        assert max_heat == [7, 4]
+
+    def test_constraint_expert_local_exchange(self):
+        current = [[[0, 1], [2, 3]]]
+        global_dep = [[[1, 0], [3, 2]]]
+        new_dep = DynamicEplb.constraint_expert_local_exchange(
+            current, global_dep)
+        assert new_dep == [[[0, 1], [2, 3]]]
+
+    def test_compute_balanced_pack_redundancy_normal(self):
+        origin_weights = [(0, 10), (1, 20)]
+        result, boxes = DynamicEplb.compute_balanced_pack_redundancy(
+            origin_weights, 2, 1)
+        assert isinstance(result, list) and len(result) == 2
+
+    def test_compute_balanced_pack_redundancy_card0(self):
+        origin_weights = [(0, 10)]
+        with pytest.raises(RuntimeError):
+            DynamicEplb.compute_balanced_pack_redundancy(origin_weights, 0, 0)
+
+    def test_compute_balanced_pack_normal(self):
+        origin_weights = np.array([(0, 10), (1, 20)], dtype=object)
+        result, boxes = DynamicEplb.compute_balanced_pack(origin_weights, 2)
+        assert isinstance(result, list) and len(result) == 2
+
+    def test_compute_balanced_pack_card0(self):
+        origin_weights = np.array([(0, 10)], dtype=object)
+        with pytest.raises(RuntimeError):
+            DynamicEplb.compute_balanced_pack(origin_weights, 0)
+
+    def test_original_compute_balanced_pack_redundancy(self):
+        origin_weights = [(0, 5), (1, 10)]
+        result, boxes = DynamicEplb.original_compute_balanced_pack_redundancy(
+            origin_weights, 2, 1)
+        assert isinstance(result, list) and len(result) == 2
+
+    def test_rebalance_experts_normal(self):
+        expert_table = np.array([[[0, 1], [1, 0]]])
+        workload = np.array([[[2, 3], [4, 1]]])
+        policy = DynamicEplb(config=None)
+        change, priority, new_dep = policy.rebalance_experts(
+            expert_table, workload)
+        assert change in [0, 1]
+        assert isinstance(priority, np.ndarray)
+        assert isinstance(new_dep, list)
+        assert np.array(new_dep).shape == expert_table.shape
+
+    def test_rebalance_experts_exceptions(self):
+        policy = DynamicEplb(config=None)
+
+        # case1: num_original_expert != expert_num
+        expert_table = np.array([[[0, 1], [1, 0]]])
+        workload = np.array([[[2, 3], [4, 1]]])
+        with patch.object(DynamicEplb,
+                          'add_redundant',
+                          return_value=np.array([[1, 2, 3]])):
+            with pytest.raises(ValueError):
+                policy.rebalance_experts(expert_table, workload)
+
+        # case2: num_npus <= 0
+        expert_table_zero = np.array([[]])  # 1 layer, 0 NPU, 0 experts
+        workload_zero = np.array([[]])
+        with pytest.raises(ValueError):
+            policy.rebalance_experts(expert_table_zero, workload_zero)
+
+        # case3: num_npus < num_redundancy_expert
+        expert_table_small = np.array([[[0, 0]]])  # 1 layer, 1 NPU, 2 experts
+        workload_small = np.array([[[1, 1]]])
+        with patch.object(DynamicEplb, 'get_redundant_num', return_value=2):
+            with pytest.raises(ValueError):
+                policy.rebalance_experts(expert_table_small, workload_small)

tests/ut/eplb/core/policy/test_policy_dynamic_ep_v2.py

@@ -0,0 +1,99 @@
+from typing import Dict, Set
+
+import numpy as np
+import pytest
+
+from vllm_ascend.eplb.core.policy.policy_dynamic_ep_v2 import (DynamicConfig,
+                                                               DynamicEplbV2)
+
+
+@pytest.fixture
+def config():
+    return DynamicConfig()
+
+
+@pytest.fixture
+def policy(config):
+    return DynamicEplbV2(config)
+
+
+def test_safe_operations(policy):
+    # safe_divide
+    assert policy.safe_divide(10, 2) == 5
+    assert policy.safe_divide(1, 0) == 0
+
+    # safe_exact_divide
+    assert policy.safe_exact_divide(10, 3) == 3
+    assert policy.safe_exact_divide(1, 0) == 0
+
+    # safe_mod
+    assert policy.safe_mod(10, 3) == 1
+    assert policy.safe_mod(1, 0) == 0
+
+
+def test_add_redundant():
+    workload = np.array([[[1, 2], [3, 4]]])
+    placement = np.array([[[0, 1], [0, 1]]])
+    result = DynamicEplbV2.add_redundant(placement, workload, 2)
+    assert result.shape == (1, 2)
+    assert np.all(result[0] == [4, 6])  # 0:1+3, 1:2+4
+
+
+def test_get_redundant_num():
+    counts = np.array([1, 2, 1])
+    assert DynamicEplbV2.get_redundant_num(3, counts) == 1  # sum(counts-1)
+
+
+def test_calculate_max_heat_per_layer():
+    workload = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    result = DynamicEplbV2.calculate_max_heat_per_layer(workload, 2)
+    assert result == [7, 15]
+
+
+def test_calculate_initial_imbalance(policy):
+    deployment = np.array([[[0, 1], [0, 1]]])
+    workloads = np.array([[1, 1]])
+    result = policy.calculate_initial_imbalance(deployment, workloads)
+    assert isinstance(result, list)
+    assert len(result) == 1
+
+
+def test_compute_redundant_assignments(policy):
+    base_experts = [(0, 10), (1, 5)]
+    redundant, sorted_weights = policy.compute_redundant_assignments(
+        base_experts, num_redundant_experts=2, num_experts=2)
+    assert len(redundant) == 2
+    assert len(sorted_weights) == 2
+
+
+def test_prepare_expert_list():
+    base_experts = [(0, 10), (1, 5)]
+    redundant_assignments = [[2], []]
+    result = DynamicEplbV2.prepare_expert_list(base_experts,
+                                               redundant_assignments, 1)
+    assert isinstance(result, list)
+    assert len(result) == 1
+
+
+def test_non_redundant_expert_information():
+    origin_deployment = np.array([[0, 1]])
+    updated_weights = [(0, 10), (1, 5)]
+    rendun_pos: Dict[int, Set[int]] = {0: set()}
+    assignments, weights, loads, counts = DynamicEplbV2.non_redundant_expert_information(
+        origin_deployment, updated_weights, rendun_pos)
+    assert assignments[0] == [0, 1]
+    assert loads[0] == 15
+
+
+def test_recomputing_initial_weight(policy):
+    layer_workloads = [10, 5]
+    device_assignments = [[0, 1]]
+    cur_layer_workload, num_all_experts = policy.recomputing_initial_weight(
+        layer_workloads, device_assignments)
+    assert cur_layer_workload[0] == 10
+    assert num_all_experts[0] == 1
+
+
+def test_safe_divide_zero_edge_case(policy):
+    assert policy.safe_divide(0, 1) == 0
+    assert policy.safe_divide(0, 5) == 0

tests/ut/eplb/core/policy/test_policy_factor.py

@@ -0,0 +1,23 @@
+import pytest
+
+from vllm_ascend.eplb.core.policy.policy_abstract import DynamicConfig
+from vllm_ascend.eplb.core.policy.policy_dynamic_ep import DynamicEplb
+from vllm_ascend.eplb.core.policy.policy_dynamic_ep_v2 import DynamicEplbV2
+from vllm_ascend.eplb.core.policy.policy_factory import PolicyFactory
+from vllm_ascend.eplb.core.policy.policy_random import RandomLoadBalance
+
+
+@pytest.fixture
+def dummy_config():
+    return DynamicConfig()
+
+
+@pytest.mark.parametrize("policy_type, expected_class", [
+    (0, RandomLoadBalance),
+    (1, DynamicEplb),
+    (2, DynamicEplbV2),
+    (999, RandomLoadBalance),
+])
+def test_generate_policy(policy_type, expected_class, dummy_config):
+    policy_instance = PolicyFactory.generate_policy(policy_type, dummy_config)
+    assert isinstance(policy_instance, expected_class)