Files
ray-project--ray/python/ray/tests/unit/test_node_affinity_validation.py
T
2026-07-13 13:17:40 +08:00

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__]))