def test_gpu_slot_discovery(self, ray_shutdown):
ray.init(num_cpus=4, num_gpus=4)
discovery = RayHostDiscovery(
use_gpu=True, cpus_per_worker=1, gpus_per_worker=2)
mapping = discovery.find_available_hosts_and_slots()
assert len(mapping) == 1
assert list(mapping.values()) == [2]
def test_multinode_mismatch(self, monkeypatch):
def create_multi_node_mock():
host_names = ["host-1", "host-2", "host-3"]
resources = {"CPU": 8}
def create_node_entry(hostname):
return {
"NodeManagerAddress": hostname,
"Resources": resources.copy(),
"alive": True
}
return map(create_node_entry, host_names)
monkeypatch.setattr(ray, "nodes", create_multi_node_mock)
discovery = RayHostDiscovery(use_gpu=True, cpus_per_worker=1)
mapping = discovery.find_available_hosts_and_slots()
assert sum(mapping.values()) == 0