wehub-resource-sync
398 lines
12 KiB
Python
398 lines
12 KiB
Python
import json
|
|
import platform
|
|
import random
|
|
import re
|
|
import shutil
|
|
import sys
|
|
import time
|
|
import zlib
|
|
from collections import defaultdict
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.cluster_utils import Cluster, cluster_not_supported
|
|
from ray.tests.test_object_spilling import is_dir_empty
|
|
|
|
# Note: Disk write speed can be as low as 6 MiB/s in AWS Mac instances, so we have to
|
|
# increase the timeout.
|
|
pytestmark = [pytest.mark.timeout(900 if platform.system() == "Darwin" else 180)]
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() in ["Windows"], reason="Failing on Windows.")
|
|
def test_multiple_directories(tmp_path, shutdown_only):
|
|
num_dirs = 3
|
|
temp_dirs = []
|
|
for i in range(num_dirs):
|
|
temp_folder = tmp_path / f"spill_{i}"
|
|
temp_folder.mkdir()
|
|
temp_dirs.append(temp_folder)
|
|
|
|
# Limit our object store to 75 MiB of memory.
|
|
min_spilling_size = 0
|
|
object_spilling_config = json.dumps(
|
|
{
|
|
"type": "filesystem",
|
|
"params": {"directory_path": [str(directory) for directory in temp_dirs]},
|
|
}
|
|
)
|
|
ray_context = ray.init(
|
|
object_store_memory=75 * 1024 * 1024,
|
|
_system_config={
|
|
"max_io_workers": 5,
|
|
"object_store_full_delay_ms": 100,
|
|
"object_spilling_config": object_spilling_config,
|
|
"min_spilling_size": min_spilling_size,
|
|
},
|
|
)
|
|
|
|
arr = np.ones(74 * 1024 * 1024, dtype=np.uint8) # 74MB.
|
|
object_refs = []
|
|
# Now the storage is full.
|
|
object_refs.append(ray.put(arr))
|
|
|
|
num_object_spilled = 20
|
|
for _ in range(num_object_spilled):
|
|
object_refs.append(ray.put(arr))
|
|
|
|
num_files = defaultdict(int)
|
|
for temp_dir in temp_dirs:
|
|
# Under temp_dir there are spilled_objects_dir(s) with name pattern
|
|
# "ray_spilled_objects[_<node_id>]", each containing spilled object files.
|
|
for spilled_objects_dir in temp_folder.iterdir():
|
|
for path in spilled_objects_dir.iterdir():
|
|
num_files[str(temp_folder)] += 1
|
|
|
|
for ref in object_refs:
|
|
assert np.array_equal(ray.get(ref), arr)
|
|
|
|
print("Check distribution...")
|
|
min_count = 5
|
|
is_distributed = [n_files >= min_count for n_files in num_files.values()]
|
|
assert all(is_distributed)
|
|
|
|
print("Check deletion...")
|
|
# Empty object refs.
|
|
object_refs = []
|
|
# Add a new small object so that the last entry is evicted and we don't
|
|
# exceed the spill threshold.
|
|
ref = ray.put(np.ones(5 * 1024 * 1024, dtype=np.uint8))
|
|
for temp_dir in temp_dirs:
|
|
temp_folder = temp_dir
|
|
wait_for_condition(lambda: is_dir_empty(temp_folder, ray_context["node_id"]))
|
|
|
|
# Now kill ray and see all directories are deleted.
|
|
print("Check directories are deleted...")
|
|
ray.shutdown()
|
|
for temp_dir in temp_dirs:
|
|
wait_for_condition(lambda: is_dir_empty(temp_dir, ray_context["node_id"]))
|
|
|
|
|
|
def _check_spilled(num_objects_spilled=0):
|
|
def ok():
|
|
s = ray._private.internal_api.memory_summary(stats_only=True)
|
|
if num_objects_spilled == 0:
|
|
return "Spilled " not in s
|
|
|
|
m = re.search(r"Spilled (\d+) MiB, (\d+) objects", s)
|
|
if m is not None:
|
|
actual_num_objects = int(m.group(2))
|
|
return actual_num_objects >= num_objects_spilled
|
|
|
|
return False
|
|
|
|
wait_for_condition(ok, timeout=90, retry_interval_ms=5000)
|
|
|
|
|
|
def _test_object_spilling_threshold(thres, num_objects, num_objects_spilled):
|
|
try:
|
|
ray.init(
|
|
object_store_memory=2_200_000_000,
|
|
_system_config={"object_spilling_threshold": thres} if thres else {},
|
|
)
|
|
objs = []
|
|
for _ in range(num_objects):
|
|
objs.append(ray.put(np.empty(200_000_000, dtype=np.uint8)))
|
|
if num_objects_spilled == 0:
|
|
time.sleep(10) # Wait for spilling to happen
|
|
_check_spilled(num_objects_spilled)
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() != "Linux", reason="Failing on Windows/macOS.")
|
|
def test_object_spilling_threshold_default():
|
|
_test_object_spilling_threshold(None, 10, 5)
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() != "Linux", reason="Failing on Windows/macOS.")
|
|
def test_object_spilling_threshold_1_0():
|
|
_test_object_spilling_threshold(1.0, 10, 0)
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() != "Linux", reason="Failing on Windows/macOS.")
|
|
def test_object_spilling_threshold_0_1():
|
|
_test_object_spilling_threshold(0.1, 10, 5)
|
|
|
|
|
|
def test_partial_retval_allocation(ray_start_cluster_enabled):
|
|
cluster = ray_start_cluster_enabled
|
|
cluster.add_node(object_store_memory=100 * 1024 * 1024)
|
|
ray.init(cluster.address)
|
|
|
|
@ray.remote(num_returns=4)
|
|
def f():
|
|
return [np.zeros(50 * 1024 * 1024, dtype=np.uint8) for _ in range(4)]
|
|
|
|
ret = f.remote()
|
|
for obj in ret:
|
|
obj = ray.get(obj)
|
|
print(obj.size)
|
|
|
|
|
|
def test_pull_spilled_object(
|
|
ray_start_cluster_enabled, multi_node_object_spilling_config, shutdown_only
|
|
):
|
|
cluster = ray_start_cluster_enabled
|
|
object_spilling_config, _ = multi_node_object_spilling_config
|
|
|
|
# Head node.
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"custom": 0},
|
|
object_store_memory=75 * 1024 * 1024,
|
|
_system_config={
|
|
"max_io_workers": 2,
|
|
"min_spilling_size": 1 * 1024 * 1024,
|
|
"automatic_object_spilling_enabled": True,
|
|
"object_store_full_delay_ms": 100,
|
|
"object_spilling_config": object_spilling_config,
|
|
},
|
|
)
|
|
ray.init(cluster.address)
|
|
|
|
# add 1 worker node
|
|
cluster.add_node(
|
|
num_cpus=1, resources={"custom": 1}, object_store_memory=75 * 1024 * 1024
|
|
)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_cpus=1, resources={"custom": 1})
|
|
def create_objects():
|
|
results = []
|
|
for size in range(5):
|
|
arr = np.random.rand(size * 1024 * 1024)
|
|
hash_value = zlib.crc32(arr.tobytes())
|
|
results.append([ray.put(arr), hash_value])
|
|
# ensure the objects are spilled
|
|
arr = np.random.rand(5 * 1024 * 1024)
|
|
ray.get(ray.put(arr))
|
|
ray.get(ray.put(arr))
|
|
return results
|
|
|
|
@ray.remote(num_cpus=1, resources={"custom": 0})
|
|
def get_object(arr):
|
|
return zlib.crc32(arr.tobytes())
|
|
|
|
results = ray.get(create_objects.remote())
|
|
for value_ref, hash_value in results:
|
|
hash_value1 = ray.get(get_object.remote(value_ref))
|
|
assert hash_value == hash_value1
|
|
|
|
|
|
# TODO(chenshen): fix error handling when spilled file
|
|
# missing/corrupted
|
|
@pytest.mark.skipif(True, reason="Currently hangs.")
|
|
def test_pull_spilled_object_failure(object_spilling_config, ray_start_cluster):
|
|
object_spilling_config, temp_folder = object_spilling_config
|
|
cluster = ray_start_cluster
|
|
|
|
# Head node.
|
|
cluster.add_node(
|
|
num_cpus=1,
|
|
resources={"custom": 0},
|
|
object_store_memory=75 * 1024 * 1024,
|
|
_system_config={
|
|
"max_io_workers": 2,
|
|
"min_spilling_size": 1 * 1024 * 1024,
|
|
"automatic_object_spilling_enabled": True,
|
|
"object_store_full_delay_ms": 100,
|
|
"object_spilling_config": object_spilling_config,
|
|
},
|
|
)
|
|
ray.init(cluster.address)
|
|
|
|
# add 1 worker node
|
|
cluster.add_node(
|
|
num_cpus=1, resources={"custom": 1}, object_store_memory=75 * 1024 * 1024
|
|
)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote(num_cpus=1, resources={"custom": 1})
|
|
def create_objects():
|
|
arr = np.random.rand(5 * 1024 * 1024)
|
|
hash_value = zlib.crc32(arr.tobytes())
|
|
results = [ray.put(arr), hash_value]
|
|
# ensure the objects are spilled
|
|
arr = np.random.rand(5 * 1024 * 1024)
|
|
ray.get(ray.put(arr))
|
|
ray.get(ray.put(arr))
|
|
return results
|
|
|
|
@ray.remote(num_cpus=1, resources={"custom": 0})
|
|
def get_object(arr):
|
|
return zlib.crc32(arr.tobytes())
|
|
|
|
[ref, hash_value] = ray.get(create_objects.remote())
|
|
|
|
# remove spilled file
|
|
shutil.rmtree(temp_folder)
|
|
|
|
hash_value1 = ray.get(get_object.remote(ref))
|
|
assert hash_value == hash_value1
|
|
|
|
|
|
@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
|
|
def test_spill_dir_cleanup_on_node_removal(fs_only_object_spilling_config):
|
|
object_spilling_config, temp_folder = fs_only_object_spilling_config
|
|
cluster = Cluster()
|
|
cluster.add_node(
|
|
num_cpus=0,
|
|
object_store_memory=75 * 1024 * 1024,
|
|
_system_config={"object_spilling_config": object_spilling_config},
|
|
)
|
|
ray.init(address=cluster.address)
|
|
node2 = cluster.add_node(num_cpus=1, object_store_memory=75 * 1024 * 1024)
|
|
|
|
# This task will run on node 2 because node 1 has no CPU resource
|
|
@ray.remote(num_cpus=1)
|
|
def run_workload():
|
|
ids = []
|
|
for _ in range(2):
|
|
arr = np.random.rand(5 * 1024 * 1024) # 40 MB
|
|
ids.append(ray.put(arr))
|
|
return ids
|
|
|
|
ids = ray.get(run_workload.remote())
|
|
node2_id = node2.node_id
|
|
assert not is_dir_empty(temp_folder, node2_id)
|
|
|
|
# Kill node 2
|
|
cluster.remove_node(node2)
|
|
|
|
# Verify that the spill folder is cleaned up upon node removal
|
|
assert is_dir_empty(temp_folder, node2_id)
|
|
|
|
# We hold the object refs to prevent them from being deleted
|
|
# due to out of scope.
|
|
del ids
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
def test_spill_deadlock(object_spilling_config, shutdown_only):
|
|
object_spilling_config, _ = object_spilling_config
|
|
# Limit our object store to 75 MiB of memory.
|
|
ray.init(
|
|
object_store_memory=75 * 1024 * 1024,
|
|
_system_config={
|
|
"max_io_workers": 1,
|
|
"automatic_object_spilling_enabled": True,
|
|
"object_store_full_delay_ms": 100,
|
|
"object_spilling_config": object_spilling_config,
|
|
"min_spilling_size": 0,
|
|
},
|
|
)
|
|
arr = np.random.rand(1024 * 1024) # 8 MB data
|
|
replay_buffer = []
|
|
|
|
# Create objects of more than 400 MiB.
|
|
for _ in range(50):
|
|
ref = None
|
|
while ref is None:
|
|
ref = ray.put(arr)
|
|
replay_buffer.append(ref)
|
|
# This is doing random sampling with 50% prob.
|
|
if random.randint(0, 9) < 5:
|
|
for _ in range(5):
|
|
ref = random.choice(replay_buffer)
|
|
sample = ray.get(ref, timeout=None)
|
|
assert np.array_equal(sample, arr)
|
|
|
|
|
|
def test_spill_reconstruction_errors(ray_start_cluster, object_spilling_config):
|
|
config = {
|
|
"health_check_failure_threshold": 10,
|
|
"health_check_period_ms": 100,
|
|
"health_check_initial_delay_ms": 0,
|
|
"max_direct_call_object_size": 100,
|
|
"task_retry_delay_ms": 100,
|
|
"object_timeout_milliseconds": 200,
|
|
}
|
|
cluster = ray_start_cluster
|
|
# Head node with no resources.
|
|
cluster.add_node(num_cpus=0, _system_config=config, object_store_memory=10**8)
|
|
ray.init(address=cluster.address)
|
|
# Node to place the initial object.
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
|
|
@ray.remote
|
|
def put():
|
|
return np.zeros(10**5, dtype=np.uint8)
|
|
|
|
@ray.remote
|
|
def check(x):
|
|
return
|
|
|
|
ref = put.remote()
|
|
for _ in range(4):
|
|
ray.get(check.remote(ref))
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
|
|
# All reconstruction attempts used up. The object's value should now be an
|
|
# error in the local store.
|
|
# Force object spilling and check that it can complete.
|
|
xs = []
|
|
for _ in range(20):
|
|
xs.append(ray.put(np.zeros(10**7, dtype=np.uint8)))
|
|
for x in xs:
|
|
ray.get(x, timeout=10)
|
|
|
|
with pytest.raises(ray.exceptions.ObjectLostError):
|
|
ray.get(ref)
|
|
|
|
|
|
def test_evict_secondary_copies_before_spill(ray_start_cluster, object_spilling_config):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
ray.init(address=cluster.address)
|
|
for _ in range(3):
|
|
cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
wait_for_condition(lambda: ray.cluster_resources()["CPU"] >= 4)
|
|
|
|
# Spread data onto all nodes.
|
|
@ray.remote
|
|
def gen():
|
|
time.sleep(0.5)
|
|
return np.ones(10 * 1024 * 1024, dtype=np.uint8)
|
|
|
|
refs = [
|
|
gen.options(scheduling_strategy="SPREAD").remote() for _ in range(16)
|
|
] # 160MiB
|
|
|
|
# Iterate over the data on the worker nodes from the head node.
|
|
for i in range(10):
|
|
for j, r in enumerate(refs):
|
|
print("Iteration", i, j)
|
|
ray.get(r)
|
|
|
|
_check_spilled()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|