49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
import re
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
from ray._raylet import NodeID
|
|
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
|
|
|
|
|
def assert_invalid_node_id(node_id_value):
|
|
node_id_str = (
|
|
node_id_value if isinstance(node_id_value, str) else node_id_value.hex()
|
|
)
|
|
expected_msg = re.escape(
|
|
f"Invalid node_id '{node_id_str}'. Node ID must be a valid "
|
|
"hex string. To get a list of all nodes and their IDs in your cluster, "
|
|
"use ray.nodes(). See https://docs.ray.io/en/latest/ray-core/miscellaneous.html"
|
|
"#node-information for more details."
|
|
)
|
|
with pytest.raises(ValueError, match=expected_msg):
|
|
NodeAffinitySchedulingStrategy(node_id=node_id_value, soft=False)
|
|
|
|
|
|
def test_node_affinity_scheduling_strategy_invalid_attributes():
|
|
valid_hex = NodeID.from_random().hex()
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="_spill_on_unavailable cannot be set when soft is False. "
|
|
"Please set soft to True to use _spill_on_unavailable.",
|
|
):
|
|
NodeAffinitySchedulingStrategy(
|
|
node_id=valid_hex, soft=False, _spill_on_unavailable=True
|
|
)
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="_fail_on_unavailable cannot be set when soft is True. "
|
|
"Please set soft to False to use _fail_on_unavailable.",
|
|
):
|
|
NodeAffinitySchedulingStrategy(
|
|
node_id=valid_hex, soft=True, _fail_on_unavailable=True
|
|
)
|
|
|
|
assert_invalid_node_id("invalid_node_id")
|
|
assert_invalid_node_id(NodeID.nil())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-vv", __file__]))
|