wehub-resource-sync
146 lines
4.0 KiB
Python
146 lines
4.0 KiB
Python
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.train.v2._internal.execution.worker_group.placement_group_handle import (
|
|
DefaultPlacementGroupHandle,
|
|
PlacementGroupHandle,
|
|
SlicePlacementGroupHandle,
|
|
)
|
|
from ray.util.placement_group import placement_group
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def ray_start():
|
|
ray.init(num_cpus=4)
|
|
yield
|
|
ray.shutdown()
|
|
|
|
|
|
# DefaultPlacementGroupHandle tests
|
|
|
|
|
|
def test_default_handle_is_placement_group_handle():
|
|
"""DefaultPlacementGroupHandle should be a PlacementGroupHandle."""
|
|
pg = placement_group([{"CPU": 1}])
|
|
handle = DefaultPlacementGroupHandle(pg)
|
|
assert isinstance(handle, PlacementGroupHandle)
|
|
handle.shutdown()
|
|
|
|
|
|
def test_default_handle_placement_group_property():
|
|
"""placement_group property should return the underlying PlacementGroup."""
|
|
pg = placement_group([{"CPU": 1}])
|
|
handle = DefaultPlacementGroupHandle(pg)
|
|
assert handle.placement_group is pg
|
|
handle.shutdown()
|
|
|
|
|
|
def test_default_handle_ready():
|
|
"""ready() should return an ObjectRef that can be waited on."""
|
|
pg = placement_group([{"CPU": 1}])
|
|
handle = DefaultPlacementGroupHandle(pg)
|
|
|
|
ready_ref = handle.ready()
|
|
# Should be able to ray.get the ready ref
|
|
ray.get(ready_ref, timeout=10)
|
|
handle.shutdown()
|
|
|
|
|
|
def test_default_handle_wait():
|
|
"""wait() should return True if the placement group is ready."""
|
|
pg = placement_group([{"CPU": 1}])
|
|
handle = DefaultPlacementGroupHandle(pg)
|
|
assert handle.wait(timeout_seconds=10)
|
|
handle.shutdown()
|
|
|
|
|
|
def test_default_handle_wait_not_found_returns_false():
|
|
"""wait() should return False if the placement group no longer exists."""
|
|
mock_pg = MagicMock()
|
|
mock_pg.wait.side_effect = Exception(
|
|
"Placement group PlacementGroupID(abc) does not exist."
|
|
)
|
|
handle = DefaultPlacementGroupHandle(mock_pg)
|
|
assert handle.wait(timeout_seconds=0) is False
|
|
|
|
|
|
def test_default_handle_shutdown():
|
|
"""shutdown() should remove the placement group."""
|
|
pg = placement_group([{"CPU": 1}])
|
|
handle = DefaultPlacementGroupHandle(pg)
|
|
|
|
# Wait for PG to be ready
|
|
ray.get(handle.ready(), timeout=10)
|
|
|
|
handle.shutdown()
|
|
|
|
|
|
# SlicePlacementGroupHandle tests
|
|
|
|
|
|
def test_slice_handle_is_placement_group_handle():
|
|
"""SlicePlacementGroupHandle should be a PlacementGroupHandle."""
|
|
mock_spg = MagicMock()
|
|
mock_pg = MagicMock()
|
|
mock_spg.placement_group = mock_pg
|
|
|
|
handle = SlicePlacementGroupHandle(mock_spg)
|
|
assert isinstance(handle, PlacementGroupHandle)
|
|
|
|
|
|
def test_slice_handle_placement_group_property():
|
|
"""placement_group property should return the underlying PlacementGroup."""
|
|
mock_spg = MagicMock()
|
|
mock_pg = MagicMock()
|
|
mock_spg.placement_group = mock_pg
|
|
|
|
handle = SlicePlacementGroupHandle(mock_spg)
|
|
assert handle.placement_group is mock_pg
|
|
|
|
|
|
def test_slice_handle_ready():
|
|
"""ready() should delegate to the underlying PlacementGroup."""
|
|
mock_spg = MagicMock()
|
|
mock_pg = MagicMock()
|
|
mock_ready_ref = MagicMock()
|
|
mock_pg.ready.return_value = mock_ready_ref
|
|
mock_spg.placement_group = mock_pg
|
|
|
|
handle = SlicePlacementGroupHandle(mock_spg)
|
|
result = handle.ready()
|
|
|
|
mock_pg.ready.assert_called_once()
|
|
assert result is mock_ready_ref
|
|
|
|
|
|
def test_slice_handle_wait():
|
|
"""wait() should delegate to the underlying PlacementGroup."""
|
|
mock_spg = MagicMock()
|
|
mock_pg = MagicMock()
|
|
mock_pg.wait.return_value = True
|
|
mock_spg.placement_group = mock_pg
|
|
|
|
handle = SlicePlacementGroupHandle(mock_spg)
|
|
assert handle.wait(timeout_seconds=10)
|
|
mock_pg.wait.assert_called_once()
|
|
|
|
|
|
def test_slice_handle_shutdown():
|
|
"""shutdown() should call shutdown on the SlicePlacementGroup."""
|
|
mock_spg = MagicMock()
|
|
mock_pg = MagicMock()
|
|
mock_spg.placement_group = mock_pg
|
|
|
|
handle = SlicePlacementGroupHandle(mock_spg)
|
|
handle.shutdown()
|
|
|
|
mock_spg.shutdown.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|