wehub-resource-sync
1278 lines
52 KiB
Python
1278 lines
52 KiB
Python
import pytest
|
|
|
|
from ray.serve._private.common import DeploymentID, ReplicaID
|
|
from ray.serve._private.deployment_state import DeploymentRankManager
|
|
from ray.serve.schema import ReplicaRank
|
|
|
|
|
|
@pytest.fixture
|
|
def rank_manager() -> DeploymentRankManager:
|
|
"""Fixture providing a fresh DeploymentRankManager instance for each test."""
|
|
return DeploymentRankManager()
|
|
|
|
|
|
class MockDeploymentReplica:
|
|
"""Mock replica for testing without heavy dependencies."""
|
|
|
|
def __init__(
|
|
self,
|
|
replica_id: str,
|
|
deployment_name: str = "test_deployment",
|
|
app_name: str = "test_app",
|
|
):
|
|
self.replica_id = ReplicaID(
|
|
unique_id=replica_id,
|
|
deployment_id=DeploymentID(name=deployment_name, app_name=app_name),
|
|
)
|
|
|
|
def __str__(self):
|
|
return f"MockDeploymentReplica(replica_id={self.replica_id})"
|
|
|
|
|
|
class TestDeploymentRankManager:
|
|
"""Test cases for DeploymentRankManager."""
|
|
|
|
def test_init(self, rank_manager):
|
|
"""Test initialization creates empty state."""
|
|
assert rank_manager.get_replica_ranks_mapping() == {}
|
|
|
|
def test_assign_rank_first_replica(self, rank_manager: DeploymentRankManager):
|
|
"""Test assigning rank to first replica."""
|
|
rank = rank_manager.assign_rank("replica_1", "node_1")
|
|
assert rank.rank == 0
|
|
assert rank_manager.has_replica_rank("replica_1")
|
|
assert rank_manager.get_replica_rank("replica_1").rank == 0
|
|
|
|
def test_assign_rank_multiple_replicas(self, rank_manager: DeploymentRankManager):
|
|
"""Test assigning ranks to multiple replicas."""
|
|
rank1 = rank_manager.assign_rank("replica_1", "node_1")
|
|
rank2 = rank_manager.assign_rank("replica_2", "node_1")
|
|
rank3 = rank_manager.assign_rank("replica_3", "node_1")
|
|
|
|
assert rank1.rank == 0
|
|
assert rank2.rank == 1
|
|
assert rank3.rank == 2
|
|
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
assert len(mapping) == 3
|
|
assert mapping == {
|
|
"replica_1": ReplicaRank(rank=0, node_rank=0, local_rank=0),
|
|
"replica_2": ReplicaRank(rank=1, node_rank=0, local_rank=1),
|
|
"replica_3": ReplicaRank(rank=2, node_rank=0, local_rank=2),
|
|
}
|
|
|
|
def test_assign_rank_reuses_released_ranks(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test that released ranks are reused before assigning new ones."""
|
|
# Assign ranks to 3 replicas
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("replica_2", "node_1")
|
|
rank_manager.assign_rank("replica_3", "node_1")
|
|
|
|
# Release middle rank
|
|
rank_manager.release_rank("replica_2")
|
|
assert not rank_manager.has_replica_rank("replica_2")
|
|
|
|
# New replica should get the released rank (1)
|
|
rank = rank_manager.assign_rank("replica_4", "node_1")
|
|
assert rank.rank == 1
|
|
assert rank_manager.get_replica_rank("replica_4") == ReplicaRank(
|
|
rank=1, node_rank=0, local_rank=1
|
|
)
|
|
|
|
def test_assign_rank_duplicate_fails(self, rank_manager: DeploymentRankManager):
|
|
"""Test assigning rank to replica that already has one fails."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
with pytest.raises(RuntimeError, match="already assigned"):
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
def test_release_rank(self, rank_manager: DeploymentRankManager):
|
|
"""Test releasing a rank makes it available for reuse."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("replica_2", "node_1")
|
|
|
|
rank_manager.release_rank("replica_1")
|
|
|
|
assert not rank_manager.has_replica_rank("replica_1")
|
|
assert rank_manager.has_replica_rank("replica_2")
|
|
assert rank_manager.get_replica_rank("replica_2").rank == 1
|
|
|
|
def test_release_rank_nonexistent_replica(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test releasing rank for non-existent replica fails."""
|
|
with pytest.raises(RuntimeError, match="not assigned"):
|
|
rank_manager.release_rank("nonexistent")
|
|
|
|
def test_recover_rank_basic(self, rank_manager: DeploymentRankManager):
|
|
"""Test basic rank recovery."""
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
assert rank_manager.has_replica_rank("replica_1")
|
|
assert rank_manager.get_replica_rank("replica_1").rank == 5
|
|
|
|
def test_recover_rank_updates_next_rank(self, rank_manager: DeploymentRankManager):
|
|
"""Test that recovering a high rank updates next_rank appropriately."""
|
|
rank_manager.assign_rank("replica_1", "node_1") # Gets rank 0
|
|
rank_manager.recover_rank(
|
|
"replica_2", "node_1", ReplicaRank(rank=10, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
# New replica should get rank 11 (next available after 10)
|
|
rank = rank_manager.assign_rank("replica_3", "node_1")
|
|
assert rank.rank == 11
|
|
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
assert mapping == {
|
|
"replica_1": ReplicaRank(rank=0, node_rank=0, local_rank=0),
|
|
"replica_2": ReplicaRank(rank=10, node_rank=0, local_rank=0),
|
|
"replica_3": ReplicaRank(rank=11, node_rank=0, local_rank=1),
|
|
}
|
|
|
|
def test_recover_rank_removes_from_available(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test that recovering a rank removes it from available ranks."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("replica_2", "node_1")
|
|
rank_manager.release_rank("replica_1") # Rank 0 becomes available
|
|
|
|
# Recover rank 0 for a new replica
|
|
rank_manager.recover_rank(
|
|
"replica_3", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
# Verify replica_3 has rank 0
|
|
assert rank_manager.has_replica_rank("replica_3")
|
|
assert rank_manager.get_replica_rank("replica_3").rank == 0
|
|
|
|
# Next assigned replica should get rank 2 (not 0, which is now taken)
|
|
rank = rank_manager.assign_rank("replica_4", "node_1")
|
|
assert rank.rank == 2
|
|
|
|
def test_recover_rank_duplicate_fails(self, rank_manager: DeploymentRankManager):
|
|
"""Test recovering rank for replica that already has one fails."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
with pytest.raises(RuntimeError, match="already assigned"):
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
def test_get_replica_rank_existing(self, rank_manager: DeploymentRankManager):
|
|
"""Test getting rank for existing replica."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank = rank_manager.get_replica_rank("replica_1")
|
|
assert rank.rank == 0
|
|
|
|
def test_get_replica_rank_nonexistent_fails(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test getting rank for non-existent replica fails."""
|
|
with pytest.raises(RuntimeError, match="not assigned"):
|
|
rank_manager.get_replica_rank("nonexistent")
|
|
|
|
def test_get_replica_ranks_mapping(self, rank_manager: DeploymentRankManager):
|
|
"""Test getting copy of replica ranks mapping."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("replica_2", "node_1")
|
|
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
expected = {
|
|
"replica_1": ReplicaRank(rank=0, node_rank=0, local_rank=0),
|
|
"replica_2": ReplicaRank(rank=1, node_rank=0, local_rank=1),
|
|
}
|
|
|
|
assert mapping == expected
|
|
|
|
# Verify it's a copy by modifying it
|
|
mapping["replica_3"] = ReplicaRank(rank=2, node_rank=0, local_rank=2)
|
|
# Get a fresh mapping to verify the original wasn't changed
|
|
fresh_mapping = rank_manager.get_replica_ranks_mapping()
|
|
assert "replica_3" not in fresh_mapping
|
|
assert fresh_mapping == expected
|
|
|
|
def test_clear(self, rank_manager: DeploymentRankManager):
|
|
"""Test clearing all rank data."""
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("replica_2", "node_1")
|
|
rank_manager.release_rank("replica_1")
|
|
|
|
rank_manager.clear()
|
|
|
|
# After clearing, should have no ranks
|
|
assert rank_manager.get_replica_ranks_mapping() == {}
|
|
assert not rank_manager.has_replica_rank("replica_1")
|
|
assert not rank_manager.has_replica_rank("replica_2")
|
|
|
|
# Should be able to assign from 0 again
|
|
rank = rank_manager.assign_rank("replica_3", "node_1")
|
|
assert rank.rank == 0
|
|
|
|
def test_check_rank_consistency_empty_replicas(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test consistency check with no active replicas."""
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally([])
|
|
assert result == []
|
|
|
|
def test_check_rank_consistency_contiguous_ranks(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test consistency check with contiguous ranks (no reassignment needed)."""
|
|
# Set up contiguous ranks
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
replica2 = MockDeploymentReplica("replica_2")
|
|
replica3 = MockDeploymentReplica("replica_3")
|
|
|
|
rank_manager.assign_rank("replica_1", "node_1") # rank 0
|
|
rank_manager.assign_rank("replica_2", "node_1") # rank 1
|
|
rank_manager.assign_rank("replica_3", "node_1") # rank 2
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3]
|
|
)
|
|
|
|
assert result == []
|
|
|
|
def test_check_rank_consistency_non_contiguous_ranks(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test consistency check with non-contiguous ranks (reassignment needed)."""
|
|
# Set up non-contiguous ranks (simulate a replica being removed)
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
replica2 = MockDeploymentReplica("replica_2")
|
|
replica3 = MockDeploymentReplica("replica_3")
|
|
|
|
# Manually assign non-contiguous ranks using recover_rank
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"replica_2", "node_1", ReplicaRank(rank=2, node_rank=0, local_rank=1)
|
|
) # Gap at rank 1
|
|
rank_manager.recover_rank(
|
|
"replica_3", "node_1", ReplicaRank(rank=3, node_rank=0, local_rank=2)
|
|
)
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3]
|
|
)
|
|
|
|
# Should reassign some replicas to make ranks contiguous
|
|
assert len(result) > 0
|
|
|
|
# After reassignment, ranks should be contiguous [0, 1, 2]
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
final_ranks = sorted([r.rank for r in mapping.values()])
|
|
expected_ranks = [0, 1, 2]
|
|
assert final_ranks == expected_ranks
|
|
|
|
def test_minimal_reassignment_keeps_existing_when_possible(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test that minimal reassignment keeps existing ranks when possible."""
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
replica2 = MockDeploymentReplica("replica_2")
|
|
replica3 = MockDeploymentReplica("replica_3")
|
|
replica4 = MockDeploymentReplica("replica_4")
|
|
|
|
# Set up ranks: 0, 2, 5, 7 (non-contiguous) using recover_rank
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
) # Should keep this
|
|
rank_manager.recover_rank(
|
|
"replica_2", "node_1", ReplicaRank(rank=2, node_rank=0, local_rank=1)
|
|
) # Should keep this
|
|
rank_manager.recover_rank(
|
|
"replica_3", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=2)
|
|
) # Should be reassigned to 1
|
|
rank_manager.recover_rank(
|
|
"replica_4", "node_1", ReplicaRank(rank=7, node_rank=0, local_rank=3)
|
|
) # Should be reassigned to 3
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3, replica4]
|
|
)
|
|
|
|
# Verify minimal reassignment
|
|
assert len(result) == 2 # Only 2 replicas should be reassigned
|
|
reassigned_ids = {r.replica_id.unique_id for r in result}
|
|
assert reassigned_ids == {"replica_3", "replica_4"}
|
|
|
|
# Verify final ranks are contiguous
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
final_ranks = sorted([r.rank for r in mapping.values()])
|
|
assert final_ranks == [0, 1, 2, 3]
|
|
|
|
# Verify that replica_1 and replica_2 kept their original ranks
|
|
assert rank_manager.get_replica_rank("replica_1").rank == 0
|
|
assert rank_manager.get_replica_rank("replica_2").rank == 2
|
|
|
|
def test_check_rank_consistency_unranked_replicas_fails(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test consistency check fails when active replicas have no ranks."""
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
|
|
with pytest.raises(RuntimeError, match="Rank system is in an invalid state"):
|
|
rank_manager.check_rank_consistency_and_reassign_minimally([replica1])
|
|
|
|
def test_check_rank_consistency_stale_ranks_fails(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test consistency check fails when there are stale ranks."""
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
|
|
# Set up stale rank (replica not in active list)
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("stale_replica", "node_1")
|
|
|
|
with pytest.raises(RuntimeError, match="Rank system is in an invalid state"):
|
|
rank_manager.check_rank_consistency_and_reassign_minimally([replica1])
|
|
|
|
def test_check_rank_consistency_duplicate_ranks_fails(
|
|
self, rank_manager: DeploymentRankManager
|
|
):
|
|
"""Test consistency check fails when there are duplicate ranks."""
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
replica2 = MockDeploymentReplica("replica_2")
|
|
|
|
# Manually create duplicate ranks using recover_rank (this should never happen in normal operation)
|
|
# Note: We can only test this with duplicate global ranks, not duplicate local ranks
|
|
# since local_rank uniqueness is enforced by the underlying RankManager
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"replica_2", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=1)
|
|
) # Duplicate global rank!
|
|
|
|
with pytest.raises(RuntimeError, match="Rank system is in an invalid state"):
|
|
rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2]
|
|
)
|
|
|
|
|
|
class TestDeploymentRankManagerMultiNode:
|
|
def test_assign_rank_multiple_nodes(self):
|
|
"""Test that replicas on different nodes get correct node_rank and local_rank."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign replicas across different nodes
|
|
rank1 = rank_manager.assign_rank("replica_1", "node_1")
|
|
rank2 = rank_manager.assign_rank("replica_2", "node_2")
|
|
rank3 = rank_manager.assign_rank("replica_3", "node_1")
|
|
rank4 = rank_manager.assign_rank("replica_4", "node_2")
|
|
rank5 = rank_manager.assign_rank("replica_5", "node_3")
|
|
|
|
# Verify global ranks are sequential
|
|
assert rank1.rank == 0
|
|
assert rank2.rank == 1
|
|
assert rank3.rank == 2
|
|
assert rank4.rank == 3
|
|
assert rank5.rank == 4
|
|
|
|
# Verify node ranks are assigned correctly
|
|
assert rank1.node_rank == 0 # node_1 is first
|
|
assert rank2.node_rank == 1 # node_2 is second
|
|
assert rank3.node_rank == 0 # node_1 already has rank 0
|
|
assert rank4.node_rank == 1 # node_2 already has rank 1
|
|
assert rank5.node_rank == 2 # node_3 is third
|
|
|
|
# Verify local ranks restart per node
|
|
assert rank1.local_rank == 0 # First replica on node_1
|
|
assert rank2.local_rank == 0 # First replica on node_2
|
|
assert rank3.local_rank == 1 # Second replica on node_1
|
|
assert rank4.local_rank == 1 # Second replica on node_2
|
|
assert rank5.local_rank == 0 # First replica on node_3
|
|
|
|
def test_local_rank_independence_across_nodes(self):
|
|
"""Test that local_rank sequences are independent per node."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Add 3 replicas to node_1
|
|
r1 = rank_manager.assign_rank("n1_r1", "node_1")
|
|
r2 = rank_manager.assign_rank("n1_r2", "node_1")
|
|
r3 = rank_manager.assign_rank("n1_r3", "node_1")
|
|
|
|
# Add 2 replicas to node_2
|
|
r4 = rank_manager.assign_rank("n2_r1", "node_2")
|
|
r5 = rank_manager.assign_rank("n2_r2", "node_2")
|
|
|
|
# Verify local ranks are independent
|
|
assert r1.local_rank == 0
|
|
assert r2.local_rank == 1
|
|
assert r3.local_rank == 2
|
|
assert r4.local_rank == 0 # Restarts at 0 for node_2
|
|
assert r5.local_rank == 1
|
|
|
|
def test_release_rank_removes_node_when_last_replica(self):
|
|
"""Test that releasing the last replica on a node releases the node rank."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign replicas to two nodes
|
|
rank_manager.assign_rank("n1_r1", "node_1")
|
|
rank_manager.assign_rank("n2_r1", "node_2")
|
|
rank_manager.assign_rank("n2_r2", "node_2")
|
|
|
|
# Release the only replica on node_1
|
|
rank_manager.release_rank("n1_r1")
|
|
|
|
# Now add a replica to node_3 - it should get node_rank 0 (reused from node_1)
|
|
new_rank = rank_manager.assign_rank("n3_r1", "node_3")
|
|
assert new_rank.node_rank == 0 # Reused node rank
|
|
|
|
def test_release_rank_keeps_node_when_replicas_remain(self):
|
|
"""Test that releasing a replica doesn't release node rank if other replicas remain."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign 3 replicas to node_1
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
rank_manager.assign_rank("r2", "node_1")
|
|
rank_manager.assign_rank("r3", "node_1")
|
|
|
|
# Assign replica to node_2
|
|
node2_rank = rank_manager.assign_rank("r4", "node_2")
|
|
assert node2_rank.node_rank == 1
|
|
|
|
# Release one replica from node_1
|
|
rank_manager.release_rank("r2")
|
|
|
|
# Add another replica to node_1 - should still have node_rank 0
|
|
new_rank = rank_manager.assign_rank("r5", "node_1")
|
|
assert new_rank.node_rank == 0 # Node rank preserved
|
|
|
|
def test_node_rank_reuse_after_complete_release(self):
|
|
"""Test that node ranks are reused after all replicas are removed."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Create replicas on 3 nodes
|
|
rank_manager.assign_rank("n1_r1", "node_1")
|
|
rank_manager.assign_rank("n2_r1", "node_2")
|
|
rank_manager.assign_rank("n3_r1", "node_3")
|
|
|
|
# Verify node ranks
|
|
assert rank_manager.get_replica_rank("n1_r1").node_rank == 0
|
|
assert rank_manager.get_replica_rank("n2_r1").node_rank == 1
|
|
assert rank_manager.get_replica_rank("n3_r1").node_rank == 2
|
|
|
|
# Remove all replicas from node_2 (middle node)
|
|
rank_manager.release_rank("n2_r1")
|
|
|
|
# Add replica to a new node - should reuse node_rank 1
|
|
new_rank = rank_manager.assign_rank("n4_r1", "node_4")
|
|
assert new_rank.node_rank == 1 # Reused from released node_2
|
|
|
|
def test_local_rank_reuse_within_node(self):
|
|
"""Test that local ranks are reused within a node after release."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign 3 replicas to node_1
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
rank_manager.assign_rank("r2", "node_1")
|
|
rank_manager.assign_rank("r3", "node_1")
|
|
|
|
# Release middle replica (local_rank=1)
|
|
rank_manager.release_rank("r2")
|
|
|
|
# Add new replica to node_1 - should reuse local_rank 1
|
|
new_rank = rank_manager.assign_rank("r4", "node_1")
|
|
assert new_rank.local_rank == 1 # Reused local rank
|
|
|
|
def test_recover_rank_multiple_nodes(self):
|
|
"""Test recovering ranks for replicas on different nodes."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Recover replicas on different nodes
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_2", ReplicaRank(rank=1, node_rank=1, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=2, node_rank=0, local_rank=1)
|
|
)
|
|
|
|
# Verify all ranks are recovered correctly
|
|
assert rank_manager.get_replica_rank("r1") == ReplicaRank(
|
|
rank=0, node_rank=0, local_rank=0
|
|
)
|
|
assert rank_manager.get_replica_rank("r2") == ReplicaRank(
|
|
rank=1, node_rank=1, local_rank=0
|
|
)
|
|
assert rank_manager.get_replica_rank("r3") == ReplicaRank(
|
|
rank=2, node_rank=0, local_rank=1
|
|
)
|
|
|
|
def test_recover_rank_preserves_node_rank_when_node_exists(self):
|
|
"""Test that recovering a replica on an existing node doesn't create duplicate node ranks."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Recover first replica on node_1
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
# Recover another replica on the same node
|
|
rank_manager.recover_rank(
|
|
"r2", "node_1", ReplicaRank(rank=1, node_rank=0, local_rank=1)
|
|
)
|
|
|
|
# Both should have the same node_rank
|
|
assert rank_manager.get_replica_rank("r1").node_rank == 0
|
|
assert rank_manager.get_replica_rank("r2").node_rank == 0
|
|
|
|
def test_check_rank_consistency_across_multiple_nodes(self):
|
|
"""Test consistency checking with replicas spread across nodes."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replica1 = MockDeploymentReplica("r1")
|
|
replica2 = MockDeploymentReplica("r2")
|
|
replica3 = MockDeploymentReplica("r3")
|
|
replica4 = MockDeploymentReplica("r4")
|
|
|
|
# Set up non-contiguous global ranks across nodes
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_2", ReplicaRank(rank=5, node_rank=1, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=8, node_rank=0, local_rank=1)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r4", "node_2", ReplicaRank(rank=10, node_rank=1, local_rank=1)
|
|
)
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3, replica4]
|
|
)
|
|
|
|
# Should reassign some replicas to make global ranks contiguous
|
|
assert len(result) > 0
|
|
|
|
# Verify global ranks are now contiguous
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
global_ranks = sorted([r.rank for r in mapping.values()])
|
|
assert global_ranks == [0, 1, 2, 3]
|
|
|
|
def test_check_rank_consistency_local_ranks_per_node(self):
|
|
"""Test that local rank consistency is checked independently per node."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replica1 = MockDeploymentReplica("r1")
|
|
replica2 = MockDeploymentReplica("r2")
|
|
replica3 = MockDeploymentReplica("r3")
|
|
replica4 = MockDeploymentReplica("r4")
|
|
|
|
# Set up non-contiguous local ranks on node_1, contiguous on node_2
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_1", ReplicaRank(rank=1, node_rank=0, local_rank=5)
|
|
) # Gap in local rank
|
|
rank_manager.recover_rank(
|
|
"r3", "node_2", ReplicaRank(rank=2, node_rank=1, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r4", "node_2", ReplicaRank(rank=3, node_rank=1, local_rank=1)
|
|
)
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3, replica4]
|
|
)
|
|
|
|
# Should reassign replica on node_1 with non-contiguous local rank
|
|
assert len(result) > 0
|
|
|
|
# Verify local ranks are contiguous per node
|
|
r1_rank = rank_manager.get_replica_rank("r1")
|
|
r2_rank = rank_manager.get_replica_rank("r2")
|
|
r3_rank = rank_manager.get_replica_rank("r3")
|
|
r4_rank = rank_manager.get_replica_rank("r4")
|
|
|
|
# Node 1 local ranks should be [0, 1]
|
|
node1_local_ranks = sorted([r1_rank.local_rank, r2_rank.local_rank])
|
|
assert node1_local_ranks == [0, 1]
|
|
|
|
# Node 2 local ranks should be [0, 1]
|
|
node2_local_ranks = sorted([r3_rank.local_rank, r4_rank.local_rank])
|
|
assert node2_local_ranks == [0, 1]
|
|
|
|
def test_check_rank_consistency_node_ranks(self):
|
|
"""Test that node rank consistency is maintained."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replica1 = MockDeploymentReplica("r1")
|
|
replica2 = MockDeploymentReplica("r2")
|
|
replica3 = MockDeploymentReplica("r3")
|
|
|
|
# Set up non-contiguous node ranks (0, 2, 5)
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_2", ReplicaRank(rank=1, node_rank=2, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_3", ReplicaRank(rank=2, node_rank=5, local_rank=0)
|
|
)
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3]
|
|
)
|
|
|
|
# Should reassign replicas to make node ranks contiguous
|
|
assert len(result) > 0
|
|
|
|
# Verify node ranks are now contiguous [0, 1, 2]
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
node_ranks = sorted([r.node_rank for r in mapping.values()])
|
|
assert node_ranks == [0, 1, 2]
|
|
|
|
def test_clear_with_multiple_nodes(self):
|
|
"""Test that clear() removes all node-related state."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign replicas to multiple nodes
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
rank_manager.assign_rank("r2", "node_2")
|
|
rank_manager.assign_rank("r3", "node_3")
|
|
|
|
# Clear everything
|
|
rank_manager.clear()
|
|
|
|
# Verify all state is cleared
|
|
assert rank_manager.get_replica_ranks_mapping() == {}
|
|
|
|
# Verify we can assign fresh ranks starting from 0
|
|
new_rank = rank_manager.assign_rank("r4", "node_4")
|
|
assert new_rank.rank == 0
|
|
assert new_rank.node_rank == 0
|
|
assert new_rank.local_rank == 0
|
|
|
|
def test_get_replica_ranks_mapping_multiple_nodes(self):
|
|
"""Test that get_replica_ranks_mapping includes correct values for all nodes."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Create replicas across multiple nodes
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
rank_manager.assign_rank("r2", "node_2")
|
|
rank_manager.assign_rank("r3", "node_1")
|
|
rank_manager.assign_rank("r4", "node_3")
|
|
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
|
|
# Verify all replicas are in mapping with correct values
|
|
assert len(mapping) == 4
|
|
assert mapping["r1"] == ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
assert mapping["r2"] == ReplicaRank(rank=1, node_rank=1, local_rank=0)
|
|
assert mapping["r3"] == ReplicaRank(rank=2, node_rank=0, local_rank=1)
|
|
assert mapping["r4"] == ReplicaRank(rank=3, node_rank=2, local_rank=0)
|
|
|
|
def test_complex_multi_node_lifecycle(self):
|
|
"""Test a complex scenario with adds, releases, and consistency checks across nodes."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Phase 1: Initial deployment across 3 nodes
|
|
rank_manager.assign_rank("n1_r1", "node_1")
|
|
rank_manager.assign_rank("n1_r2", "node_1")
|
|
rank_manager.assign_rank("n2_r1", "node_2")
|
|
rank_manager.assign_rank("n3_r1", "node_3")
|
|
rank_manager.assign_rank("n3_r2", "node_3")
|
|
|
|
# Phase 2: Scale down - remove some replicas
|
|
rank_manager.release_rank("n1_r2") # Remove from node_1
|
|
rank_manager.release_rank("n2_r1") # Remove all from node_2
|
|
|
|
# Phase 3: Scale up - add replicas to new and existing nodes
|
|
rank_manager.assign_rank("n1_r3", "node_1") # Add to existing node_1
|
|
rank_manager.assign_rank("n4_r1", "node_4") # New node
|
|
|
|
# Verify state is consistent
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
assert len(mapping) == 5
|
|
|
|
# Verify node ranks - node_2 was removed, so node_4 should reuse its rank
|
|
assert mapping["n4_r1"].node_rank == 1 # Reused from node_2
|
|
|
|
# Verify local ranks per node
|
|
assert mapping["n1_r1"].local_rank == 0
|
|
assert mapping["n1_r3"].local_rank == 1 # Reused local rank
|
|
assert mapping["n3_r1"].local_rank == 0
|
|
assert mapping["n3_r2"].local_rank == 1
|
|
assert mapping["n4_r1"].local_rank == 0
|
|
|
|
def test_scaling_up_and_down_across_nodes(self):
|
|
"""Test scaling operations across multiple nodes."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Scale up: Add replicas to 4 nodes
|
|
for node_idx in range(4):
|
|
for replica_idx in range(2):
|
|
replica_id = f"n{node_idx}_r{replica_idx}"
|
|
rank_manager.assign_rank(replica_id, f"node_{node_idx}")
|
|
|
|
# Verify 8 replicas total
|
|
assert len(rank_manager.get_replica_ranks_mapping()) == 8
|
|
|
|
# Scale down: Remove all replicas from nodes 1 and 3
|
|
rank_manager.release_rank("n1_r0")
|
|
rank_manager.release_rank("n1_r1")
|
|
rank_manager.release_rank("n3_r0")
|
|
rank_manager.release_rank("n3_r1")
|
|
|
|
# Verify 4 replicas remain
|
|
assert len(rank_manager.get_replica_ranks_mapping()) == 4
|
|
|
|
# Scale up again: Add to new node
|
|
new_rank = rank_manager.assign_rank("n5_r0", "node_5")
|
|
|
|
# Should reuse a released node rank (1 or 3)
|
|
assert new_rank.node_rank in [1, 3]
|
|
|
|
def test_minimal_reassignment_preserves_node_assignments(self):
|
|
"""Test that minimal reassignment doesn't move replicas between nodes."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replica1 = MockDeploymentReplica("r1")
|
|
replica2 = MockDeploymentReplica("r2")
|
|
replica3 = MockDeploymentReplica("r3")
|
|
replica4 = MockDeploymentReplica("r4")
|
|
|
|
# Set up non-contiguous ranks across nodes
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_2", ReplicaRank(rank=5, node_rank=1, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=10, node_rank=0, local_rank=1)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r4", "node_2", ReplicaRank(rank=15, node_rank=1, local_rank=1)
|
|
)
|
|
|
|
# Perform consistency check
|
|
rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3, replica4]
|
|
)
|
|
|
|
# Verify replicas stayed on their original nodes
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
|
|
# r1 and r3 should still be on node_1 (node_rank=0)
|
|
assert mapping["r1"].node_rank == mapping["r3"].node_rank == 0
|
|
|
|
# r2 and r4 should still be on node_2 (node_rank=1)
|
|
assert mapping["r2"].node_rank == mapping["r4"].node_rank == 1
|
|
|
|
# Verify local ranks are contiguous per node
|
|
assert sorted([mapping["r1"].local_rank, mapping["r3"].local_rank]) == [0, 1]
|
|
assert sorted([mapping["r2"].local_rank, mapping["r4"].local_rank]) == [0, 1]
|
|
|
|
|
|
class TestDeploymentRankManagerEdgeCases:
|
|
"""Test edge cases and corner cases for correctness."""
|
|
|
|
def test_recover_rank_updates_next_rank_correctly(self):
|
|
"""When recovering a rank >= _next_rank, verify next assigned rank is correct."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign first replica normally - gets rank 0
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
|
|
# Recover a replica with a very high rank
|
|
rank_manager.recover_rank(
|
|
"r2", "node_1", ReplicaRank(rank=100, node_rank=0, local_rank=1)
|
|
)
|
|
|
|
# Now assign a new replica - should get rank 101 (next after 100)
|
|
new_rank = rank_manager.assign_rank("r3", "node_1")
|
|
assert new_rank.rank == 101
|
|
|
|
def test_assign_after_recovery_with_very_high_rank(self):
|
|
"""Test that assignment after recovering a very high rank works correctly."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Recover with rank 1000
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=1000, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
# Assign new replica - should get 1001
|
|
new_rank = rank_manager.assign_rank("r2", "node_1")
|
|
assert new_rank.rank == 1001
|
|
assert new_rank.local_rank == 1 # Second replica on node_1
|
|
|
|
def test_recover_multiple_high_ranks_out_of_order(self):
|
|
"""Test recovering multiple high ranks in non-sequential order."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Recover ranks in random order: 50, 10, 100
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=50, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_2", ReplicaRank(rank=10, node_rank=1, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=100, node_rank=0, local_rank=1)
|
|
)
|
|
|
|
# Next assignment should get 101 (max + 1)
|
|
new_rank = rank_manager.assign_rank("r4", "node_3")
|
|
assert new_rank.rank == 101
|
|
|
|
def test_recover_rank_removes_from_released_ranks(self):
|
|
"""Test that recovering a rank that was released removes it from available set."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign and then release rank 5
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=1)
|
|
)
|
|
rank_manager.release_rank("r2") # Rank 5 is now released
|
|
|
|
# Recover a different replica with rank 5
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=1)
|
|
)
|
|
|
|
# Assign new replica - should get rank 6, not 5 (since 5 is taken)
|
|
new_rank = rank_manager.assign_rank("r4", "node_1")
|
|
assert new_rank.rank == 6
|
|
|
|
def test_reassignment_with_all_ranks_out_of_range(self):
|
|
"""Test reassignment when all replicas have ranks outside target range."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replica1 = MockDeploymentReplica("r1")
|
|
replica2 = MockDeploymentReplica("r2")
|
|
replica3 = MockDeploymentReplica("r3")
|
|
|
|
# All replicas have ranks outside [0, 1, 2] range
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=10, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_1", ReplicaRank(rank=20, node_rank=0, local_rank=1)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=30, node_rank=0, local_rank=2)
|
|
)
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3]
|
|
)
|
|
|
|
# All replicas should be reassigned
|
|
assert len(result) == 3
|
|
|
|
# Verify final ranks are exactly [0, 1, 2]
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
final_ranks = sorted([r.rank for r in mapping.values()])
|
|
assert final_ranks == [0, 1, 2]
|
|
|
|
def test_reassignment_preserves_target_ranks_exactly(self):
|
|
"""Test that after reassignment, ranks are exactly [0, N-1]."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replicas = []
|
|
for i in range(5):
|
|
replicas.append(MockDeploymentReplica(f"r{i}"))
|
|
|
|
# Create non-contiguous ranks: 0, 3, 7, 11, 15
|
|
rank_manager.recover_rank(
|
|
"r0", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=3, node_rank=0, local_rank=1)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r2", "node_1", ReplicaRank(rank=7, node_rank=0, local_rank=2)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r3", "node_1", ReplicaRank(rank=11, node_rank=0, local_rank=3)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"r4", "node_1", ReplicaRank(rank=15, node_rank=0, local_rank=4)
|
|
)
|
|
|
|
rank_manager.check_rank_consistency_and_reassign_minimally(replicas)
|
|
|
|
# Verify ranks are exactly [0, 1, 2, 3, 4]
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
final_ranks = sorted([r.rank for r in mapping.values()])
|
|
assert final_ranks == [0, 1, 2, 3, 4]
|
|
|
|
# Verify no duplicates
|
|
assert len(final_ranks) == len(set(final_ranks))
|
|
|
|
def test_multiple_sequential_releases_reuse_in_order(self):
|
|
"""Test releasing multiple ranks in sequence maintains correct state."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign ranks 0-5
|
|
for i in range(6):
|
|
rank_manager.assign_rank(f"r{i}", "node_1")
|
|
|
|
# Release ranks 0, 2, 4
|
|
rank_manager.release_rank("r0")
|
|
rank_manager.release_rank("r2")
|
|
rank_manager.release_rank("r4")
|
|
|
|
# Assign new replicas - should reuse in ascending order (min first)
|
|
new_rank1 = rank_manager.assign_rank("r6", "node_1")
|
|
new_rank2 = rank_manager.assign_rank("r7", "node_1")
|
|
new_rank3 = rank_manager.assign_rank("r8", "node_1")
|
|
|
|
# Should reuse 0, 2, 4 in that order
|
|
assert new_rank1.rank == 0
|
|
assert new_rank2.rank == 2
|
|
assert new_rank3.rank == 4
|
|
|
|
def test_interleaved_assign_release_recover(self):
|
|
"""Test complex sequence of operations maintains consistency."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Complex sequence
|
|
rank_manager.assign_rank("r1", "node_1") # rank 0
|
|
rank_manager.assign_rank("r2", "node_1") # rank 1
|
|
rank_manager.release_rank("r1") # release 0
|
|
rank_manager.recover_rank(
|
|
"r3", "node_2", ReplicaRank(rank=5, node_rank=1, local_rank=0)
|
|
)
|
|
rank_manager.assign_rank("r4", "node_1") # should get 0 (reused)
|
|
rank_manager.release_rank("r2") # release 1
|
|
rank_manager.assign_rank("r5", "node_2") # should get 1 (reused)
|
|
|
|
# Verify final state
|
|
assert rank_manager.get_replica_rank("r4").rank == 0
|
|
assert rank_manager.get_replica_rank("r5").rank == 1
|
|
assert rank_manager.get_replica_rank("r3").rank == 5
|
|
|
|
def test_release_all_then_reassign_all(self):
|
|
"""Test releasing all replicas then reassigning maintains correct state."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Assign replicas across nodes
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
rank_manager.assign_rank("r2", "node_2")
|
|
rank_manager.assign_rank("r3", "node_1")
|
|
|
|
# Release all
|
|
rank_manager.release_rank("r1")
|
|
rank_manager.release_rank("r2")
|
|
rank_manager.release_rank("r3")
|
|
|
|
# Verify mapping is empty
|
|
assert rank_manager.get_replica_ranks_mapping() == {}
|
|
|
|
# Reassign new replicas - should reuse ranks 0, 1, 2
|
|
new_rank1 = rank_manager.assign_rank("r4", "node_3")
|
|
new_rank2 = rank_manager.assign_rank("r5", "node_3")
|
|
new_rank3 = rank_manager.assign_rank("r6", "node_3")
|
|
|
|
assert new_rank1.rank == 0
|
|
assert new_rank2.rank == 1
|
|
assert new_rank3.rank == 2
|
|
|
|
def test_recover_rank_with_same_replica_different_rank(self):
|
|
"""Test that recovering the same replica_id twice raises an error."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# First recovery
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
# Try to recover same replica_id again - should raise error
|
|
with pytest.raises(RuntimeError, match="already assigned"):
|
|
rank_manager.recover_rank(
|
|
"r1", "node_1", ReplicaRank(rank=1, node_rank=0, local_rank=1)
|
|
)
|
|
|
|
def test_large_scale_rank_management(self):
|
|
"""Test with many nodes and many replicas per node."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
num_nodes = 50
|
|
replicas_per_node = 10
|
|
total_replicas = num_nodes * replicas_per_node
|
|
|
|
# Assign replicas
|
|
replica_ids = []
|
|
for node_idx in range(num_nodes):
|
|
for replica_idx in range(replicas_per_node):
|
|
replica_id = f"n{node_idx}_r{replica_idx}"
|
|
replica_ids.append(replica_id)
|
|
rank_manager.assign_rank(replica_id, f"node_{node_idx}")
|
|
|
|
# Verify total count
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
assert len(mapping) == total_replicas
|
|
|
|
# Verify global ranks are contiguous [0, total-1]
|
|
global_ranks = sorted([r.rank for r in mapping.values()])
|
|
assert global_ranks == list(range(total_replicas))
|
|
|
|
# Verify node ranks are contiguous [0, num_nodes-1]
|
|
node_ranks = sorted({r.node_rank for r in mapping.values()})
|
|
assert node_ranks == list(range(num_nodes))
|
|
|
|
# Verify local ranks per node
|
|
for node_idx in range(num_nodes):
|
|
node_replicas = [
|
|
rid for rid in replica_ids if rid.startswith(f"n{node_idx}_")
|
|
]
|
|
local_ranks = sorted([mapping[rid].local_rank for rid in node_replicas])
|
|
assert local_ranks == list(range(replicas_per_node))
|
|
|
|
def test_consistency_check_with_released_ranks_in_system(self):
|
|
"""Test consistency check works correctly when released_ranks exist."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
replica1 = MockDeploymentReplica("r1")
|
|
replica2 = MockDeploymentReplica("r2")
|
|
replica3 = MockDeploymentReplica("r3")
|
|
|
|
# Assign 5 replicas
|
|
rank_manager.assign_rank("r1", "node_1") # 0
|
|
rank_manager.assign_rank("r2", "node_1") # 1
|
|
rank_manager.assign_rank("r3", "node_1") # 2
|
|
rank_manager.assign_rank("r4", "node_1") # 3
|
|
rank_manager.assign_rank("r5", "node_1") # 4
|
|
|
|
# Release two (creating gaps)
|
|
rank_manager.release_rank("r4") # Release 3
|
|
rank_manager.release_rank("r5") # Release 4
|
|
|
|
# Check consistency with only remaining 3 replicas
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2, replica3]
|
|
)
|
|
|
|
# Should not need reassignment (ranks are already 0, 1, 2)
|
|
assert result == []
|
|
|
|
# Verify ranks are correct
|
|
mapping = rank_manager.get_replica_ranks_mapping()
|
|
assert sorted([r.rank for r in mapping.values()]) == [0, 1, 2]
|
|
|
|
def test_has_replica_rank_returns_false_for_unassigned(self):
|
|
"""Test has_replica_rank returns False for various unassigned states."""
|
|
rank_manager = DeploymentRankManager()
|
|
|
|
# Completely unassigned replica
|
|
assert not rank_manager.has_replica_rank("nonexistent")
|
|
|
|
# Assign then release
|
|
rank_manager.assign_rank("r1", "node_1")
|
|
rank_manager.release_rank("r1")
|
|
assert not rank_manager.has_replica_rank("r1")
|
|
|
|
|
|
class TestDeploymentRankManagerErrorHandling:
|
|
"""Test cases for DeploymentRankManager error handling with fail_on_rank_error flag.
|
|
|
|
This test class can be easily removed in the future when the error handling
|
|
feature is no longer needed.
|
|
"""
|
|
|
|
def test_assign_rank_error_with_fail_on_rank_error_true(self):
|
|
"""Test that assign_rank raises exception when fail_on_rank_error=True."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=True)
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
# Should raise RuntimeError for duplicate assignment
|
|
with pytest.raises(RuntimeError, match="already assigned"):
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
def test_assign_rank_error_with_fail_on_rank_error_false(self):
|
|
"""Test that assign_rank returns safe default when fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
# Should return safe default (ReplicaRank(rank=0)) instead of raising
|
|
result = rank_manager.assign_rank("replica_1", "node_1")
|
|
assert result is not None
|
|
assert isinstance(result, ReplicaRank)
|
|
assert result.rank == 0
|
|
|
|
def test_release_rank_error_with_fail_on_rank_error_true(self):
|
|
"""Test that release_rank raises exception when fail_on_rank_error=True."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=True)
|
|
|
|
# Should raise RuntimeError for non-existent replica
|
|
with pytest.raises(RuntimeError, match="not assigned"):
|
|
rank_manager.release_rank("nonexistent")
|
|
|
|
def test_release_rank_error_with_fail_on_rank_error_false(self):
|
|
"""Test that release_rank returns safe default when fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
|
|
# Should return None instead of raising
|
|
result = rank_manager.release_rank("nonexistent")
|
|
assert result is None
|
|
|
|
def test_recover_rank_error_with_fail_on_rank_error_true(self):
|
|
"""Test that recover_rank raises exception when fail_on_rank_error=True."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=True)
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
# Should raise RuntimeError for duplicate recovery
|
|
with pytest.raises(RuntimeError, match="already assigned"):
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
def test_recover_rank_error_with_fail_on_rank_error_false(self):
|
|
"""Test that recover_rank returns safe default when fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
|
|
# Should return None instead of raising
|
|
result = rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=0)
|
|
)
|
|
assert result is None
|
|
|
|
def test_get_replica_rank_error_with_fail_on_rank_error_true(self):
|
|
"""Test that get_replica_rank raises exception when fail_on_rank_error=True."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=True)
|
|
|
|
# Should raise RuntimeError for non-existent replica
|
|
with pytest.raises(RuntimeError, match="not assigned"):
|
|
rank_manager.get_replica_rank("nonexistent")
|
|
|
|
def test_get_replica_rank_error_with_fail_on_rank_error_false(self):
|
|
"""Test that get_replica_rank returns safe default when fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
|
|
# Should return safe default (ReplicaRank(rank=0)) instead of raising
|
|
result = rank_manager.get_replica_rank("nonexistent")
|
|
assert result is not None
|
|
assert isinstance(result, ReplicaRank)
|
|
assert result.rank == 0
|
|
|
|
def test_check_rank_consistency_error_with_fail_on_rank_error_true(self):
|
|
"""Test that check_rank_consistency raises exception when fail_on_rank_error=True."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=True)
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
|
|
# Set up invalid state: active replica without rank
|
|
with pytest.raises(RuntimeError, match="Rank system is in an invalid state"):
|
|
rank_manager.check_rank_consistency_and_reassign_minimally([replica1])
|
|
|
|
def test_check_rank_consistency_error_with_fail_on_rank_error_false(self):
|
|
"""Test that check_rank_consistency returns safe default when fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
|
|
# Should return empty list instead of raising
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally([replica1])
|
|
assert result == []
|
|
|
|
def test_check_rank_consistency_with_stale_ranks_error_handling(self):
|
|
"""Test check_rank_consistency with stale ranks and fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
|
|
# Set up stale rank (replica not in active list)
|
|
rank_manager.assign_rank("replica_1", "node_1")
|
|
rank_manager.assign_rank("stale_replica", "node_1")
|
|
|
|
# Should return empty list instead of raising
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally([replica1])
|
|
assert result == []
|
|
|
|
def test_check_rank_consistency_with_duplicate_ranks_error_handling(self):
|
|
"""Test check_rank_consistency with duplicate ranks and fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
replica1 = MockDeploymentReplica("replica_1")
|
|
replica2 = MockDeploymentReplica("replica_2")
|
|
|
|
# Manually create duplicate ranks
|
|
rank_manager.recover_rank(
|
|
"replica_1", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
rank_manager.recover_rank(
|
|
"replica_2", "node_1", ReplicaRank(rank=0, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
# Should return empty list instead of raising
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica1, replica2]
|
|
)
|
|
assert result == []
|
|
|
|
def test_normal_operations_work_with_fail_on_rank_error_false(self):
|
|
"""Test that normal operations still work correctly with fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
|
|
# Test normal assign
|
|
rank1 = rank_manager.assign_rank("replica_1", "node_1")
|
|
assert rank1.rank == 0
|
|
|
|
# Test normal get
|
|
rank = rank_manager.get_replica_rank("replica_1")
|
|
assert rank.rank == 0
|
|
|
|
# Test normal release
|
|
rank_manager.release_rank("replica_1")
|
|
assert not rank_manager.has_replica_rank("replica_1")
|
|
|
|
# Test normal recover
|
|
rank_manager.recover_rank(
|
|
"replica_2", "node_1", ReplicaRank(rank=5, node_rank=0, local_rank=0)
|
|
)
|
|
|
|
assert rank_manager.get_replica_rank("replica_2").rank == 5
|
|
|
|
# Test normal consistency check
|
|
replica2 = MockDeploymentReplica("replica_2")
|
|
replica3 = MockDeploymentReplica("replica_3")
|
|
rank_manager.assign_rank("replica_3", "node_1")
|
|
|
|
result = rank_manager.check_rank_consistency_and_reassign_minimally(
|
|
[replica2, replica3]
|
|
)
|
|
# Should reassign to make ranks contiguous
|
|
assert len(result) > 0
|
|
|
|
def test_multiple_errors_do_not_crash_with_fail_on_rank_error_false(self):
|
|
"""Test that multiple consecutive errors don't crash when fail_on_rank_error=False."""
|
|
rank_manager = DeploymentRankManager(fail_on_rank_error=False)
|
|
|
|
# Multiple errors in a row should all return safe defaults
|
|
result1 = rank_manager.get_replica_rank("nonexistent1")
|
|
result2 = rank_manager.get_replica_rank("nonexistent2")
|
|
result3 = rank_manager.release_rank("nonexistent3")
|
|
|
|
assert result1 is not None
|
|
assert result2 is not None
|
|
assert result3 is None
|
|
|
|
# And normal operations should still work after errors
|
|
rank = rank_manager.assign_rank("replica_1", "node_1")
|
|
assert rank.rank == 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|