wehub-resource-sync
297 lines
9.4 KiB
Python
297 lines
9.4 KiB
Python
import os
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._private.accelerators import HPUAcceleratorManager, hpu
|
|
from ray.util.placement_group import placement_group
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
|
|
|
|
def test_user_configured_more_than_visible(monkeypatch, call_ray_stop_only):
|
|
# Test more hpus are configured than visible.
|
|
monkeypatch.setenv("HABANA_VISIBLE_MODULES", "0,1,2")
|
|
with pytest.raises(ValueError):
|
|
ray.init(resources={"HPU": 4})
|
|
|
|
|
|
@patch(
|
|
"ray._private.accelerators.HPUAcceleratorManager.get_current_node_num_accelerators", # noqa: E501
|
|
return_value=4,
|
|
)
|
|
def test_auto_detected_more_than_visible(
|
|
mock_get_num_accelerators, monkeypatch, shutdown_only
|
|
):
|
|
# Test more hpus are detected than visible.
|
|
monkeypatch.setenv("HABANA_VISIBLE_MODULES", "0,1,2")
|
|
ray.init()
|
|
_ = mock_get_num_accelerators.called
|
|
assert ray.available_resources()["HPU"] == 3
|
|
|
|
|
|
@patch(
|
|
"ray._private.accelerators.HPUAcceleratorManager.get_current_node_num_accelerators", # noqa: E501
|
|
return_value=2,
|
|
)
|
|
def test_auto_detect_resources(mock_get_num_accelerators, shutdown_only):
|
|
# Test that ray node resources are filled with auto detected count.
|
|
ray.init()
|
|
_ = mock_get_num_accelerators.called
|
|
assert ray.available_resources()["HPU"] == 2
|
|
|
|
|
|
def test_get_current_process_visible_accelerator_ids():
|
|
os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] = "0,1,2"
|
|
assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() == [
|
|
"0",
|
|
"1",
|
|
"2",
|
|
] # noqa: E501
|
|
|
|
del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
|
|
assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() is None
|
|
|
|
os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] = ""
|
|
assert HPUAcceleratorManager.get_current_process_visible_accelerator_ids() == []
|
|
|
|
del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
|
|
|
|
|
|
def test_set_current_process_visible_accelerator_ids():
|
|
HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0"])
|
|
assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0"
|
|
|
|
HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1"])
|
|
assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0,1"
|
|
|
|
HPUAcceleratorManager.set_current_process_visible_accelerator_ids(["0", "1", "2"])
|
|
assert os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR] == "0,1,2"
|
|
|
|
del os.environ[hpu.HABANA_VISIBLE_DEVICES_ENV_VAR]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_config",
|
|
[
|
|
(1, False),
|
|
(0.5, True),
|
|
(3, False),
|
|
],
|
|
)
|
|
def test_validate_resource_request_quantity(test_config):
|
|
num_hpus, expect_error = test_config
|
|
|
|
if expect_error:
|
|
assert (
|
|
HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[0]
|
|
is False
|
|
)
|
|
assert (
|
|
HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[1]
|
|
is not None
|
|
)
|
|
else:
|
|
assert (
|
|
HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[0]
|
|
is True
|
|
)
|
|
assert (
|
|
HPUAcceleratorManager.validate_resource_request_quantity(num_hpus)[1]
|
|
is None
|
|
)
|
|
|
|
|
|
def test_check_accelerator_info():
|
|
|
|
if HPUAcceleratorManager.is_initialized():
|
|
assert (
|
|
"Intel-GAUDI" in HPUAcceleratorManager.get_current_node_accelerator_type()
|
|
)
|
|
else:
|
|
assert HPUAcceleratorManager.get_current_node_accelerator_type() is None
|
|
|
|
assert HPUAcceleratorManager.get_resource_name() == "HPU"
|
|
|
|
|
|
def test_decorator_args():
|
|
|
|
# This is a valid way of using the decorator.
|
|
@ray.remote(resources={"HPU": 1}) # noqa: F811
|
|
class Actor: # noqa: F811
|
|
def __init__(self):
|
|
pass
|
|
|
|
# This is a valid way of using the decorator.
|
|
@ray.remote(num_cpus=1, resources={"HPU": 1}) # noqa: F811
|
|
class Actor: # noqa: F811
|
|
def __init__(self):
|
|
pass
|
|
|
|
|
|
def test_actor_deletion_with_hpus(shutdown_only):
|
|
ray.init(num_cpus=1, resources={"HPU": 1})
|
|
|
|
# When an actor that uses an HPU exits, make sure that the HPU resources
|
|
# are released.
|
|
|
|
@ray.remote(resources={"HPU": 1})
|
|
class Actor:
|
|
def getpid(self):
|
|
return os.getpid()
|
|
|
|
for _ in range(5):
|
|
# If we can successfully create an actor, that means that enough
|
|
# HPU resources are available.
|
|
a = Actor.remote()
|
|
ray.get(a.getpid.remote())
|
|
|
|
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
|
|
def test_actor_hpus(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
num_nodes = 2
|
|
num_hpus_per_raylet = 2
|
|
for i in range(num_nodes):
|
|
cluster.add_node(num_cpus=10 * 2, resources={"HPU": num_hpus_per_raylet})
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(resources={"HPU": 1})
|
|
class Actor1:
|
|
def __init__(self):
|
|
resource_ids = ray.get_runtime_context().get_accelerator_ids()
|
|
self.hpu_ids = resource_ids.get("HPU")
|
|
|
|
def get_location_and_ids(self):
|
|
return (
|
|
ray.get_runtime_context().get_node_id(),
|
|
tuple(self.hpu_ids),
|
|
)
|
|
|
|
# Create one actor per HPU.
|
|
actors = [Actor1.remote() for _ in range(num_nodes * num_hpus_per_raylet)]
|
|
# Make sure that no two actors are assigned to the same HPU.
|
|
locations_and_ids = ray.get(
|
|
[actor.get_location_and_ids.remote() for actor in actors]
|
|
)
|
|
node_names = {location for location, hpu_id in locations_and_ids}
|
|
assert len(node_names) == num_nodes
|
|
location_actor_combinations = []
|
|
for node_name in node_names:
|
|
for hpu_id in range(num_hpus_per_raylet):
|
|
location_actor_combinations.append((node_name, (f"{hpu_id}",)))
|
|
|
|
assert set(locations_and_ids) == set(location_actor_combinations)
|
|
|
|
# Creating a new actor should fail because all of the HPUs are being
|
|
# used.
|
|
a = Actor1.remote()
|
|
ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
|
|
assert ready_ids == []
|
|
|
|
|
|
def test_actor_habana_visible_devices(shutdown_only):
|
|
"""Test user can overwrite HABANA_VISIBLE_MODULES
|
|
after the actor is created."""
|
|
ray.init(resources={"HPU": 1})
|
|
|
|
@ray.remote(resources={"HPU": 1})
|
|
class Actor:
|
|
def set_habana_visible_devices(self, habana_visible_devices):
|
|
os.environ["HABANA_VISIBLE_MODULES"] = habana_visible_devices
|
|
|
|
def get_habana_visible_devices(self):
|
|
return os.environ["HABANA_VISIBLE_MODULES"]
|
|
|
|
actor = Actor.remote()
|
|
assert ray.get(actor.get_habana_visible_devices.remote()) == "0"
|
|
ray.get(actor.set_habana_visible_devices.remote("0,1"))
|
|
assert ray.get(actor.get_habana_visible_devices.remote()) == "0,1"
|
|
|
|
|
|
def test_hpu_ids(shutdown_only):
|
|
num_hpus = 3
|
|
ray.init(num_cpus=num_hpus, resources={"HPU": num_hpus})
|
|
|
|
def get_hpu_ids(hpus_per_worker):
|
|
hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
|
|
assert len(hpu_ids) == hpus_per_worker
|
|
modules = os.environ.get("HABANA_VISIBLE_MODULES")
|
|
if modules is not None:
|
|
assert modules == ",".join([str(i) for i in hpu_ids]) # noqa
|
|
for hpu_id in hpu_ids:
|
|
assert hpu_id in [str(i) for i in range(num_hpus)]
|
|
return hpu_ids
|
|
|
|
f0 = ray.remote(resources={"HPU": 0})(lambda: get_hpu_ids(0))
|
|
f1 = ray.remote(resources={"HPU": 1})(lambda: get_hpu_ids(1))
|
|
|
|
list_of_ids = ray.get([f0.remote() for _ in range(10)])
|
|
assert list_of_ids == 10 * [[]]
|
|
ray.get([f1.remote() for _ in range(10)])
|
|
|
|
# Test that actors have HABANA_VISIBLE_MODULES set properly.
|
|
|
|
def _check_hpu_env(expected_num_hpus):
|
|
hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
|
|
assert len(hpu_ids) == expected_num_hpus
|
|
if expected_num_hpus > 0:
|
|
assert os.environ["HABANA_VISIBLE_MODULES"] == ",".join(
|
|
[str(i) for i in hpu_ids] # noqa
|
|
)
|
|
else:
|
|
assert os.environ.get("HABANA_VISIBLE_MODULES") is None
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self, num_hpus):
|
|
self.num_hpus = num_hpus
|
|
_check_hpu_env(num_hpus)
|
|
self.x = num_hpus
|
|
|
|
def test(self):
|
|
_check_hpu_env(self.num_hpus)
|
|
return self.x
|
|
|
|
a0 = Actor.remote(0)
|
|
assert ray.get(a0.test.remote()) == 0
|
|
|
|
a1 = Actor.options(resources={"HPU": 1}).remote(1)
|
|
assert ray.get(a1.test.remote()) == 1
|
|
|
|
|
|
def test_hpu_with_placement_group(shutdown_only):
|
|
num_hpus = 2
|
|
ray.init(num_cpus=1, resources={"HPU": num_hpus})
|
|
|
|
@ray.remote(resources={"HPU": num_hpus})
|
|
class HPUActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def ready(self):
|
|
hpu_ids = ray.get_runtime_context().get_accelerator_ids()["HPU"]
|
|
assert len(hpu_ids) == num_hpus
|
|
assert os.environ["HABANA_VISIBLE_MODULES"] == ",".join(
|
|
[str(i) for i in hpu_ids] # noqa
|
|
)
|
|
|
|
# Reserve a placement group of 1 bundle that reserves 1 CPU and 2 HPU.
|
|
pg = placement_group([{"CPU": 1, "HPU": num_hpus}])
|
|
|
|
# Wait until placement group is created.
|
|
ray.get(pg.ready(), timeout=10)
|
|
|
|
actor = HPUActor.options(
|
|
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
|
placement_group=pg,
|
|
)
|
|
).remote()
|
|
|
|
ray.get(actor.ready.remote(), timeout=10)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|