chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,378 @@
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import random
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.network_utils import build_address
|
||||
from ray._common.test_utils import (
|
||||
fetch_prometheus,
|
||||
wait_for_condition,
|
||||
)
|
||||
from ray._private.test_utils import check_spilled_mb
|
||||
|
||||
import psutil
|
||||
|
||||
MB = 1024 * 1024
|
||||
|
||||
# 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(1800 if platform.system() == "Darwin" else 180)]
|
||||
|
||||
|
||||
def _init_ray():
|
||||
return ray.init(
|
||||
num_cpus=2,
|
||||
object_store_memory=700e6,
|
||||
object_spilling_directory="/tmp/ray/plasma",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_fallback_when_spilling_impossible_on_put():
|
||||
try:
|
||||
address = _init_ray()
|
||||
x1 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
x1p = ray.get(x1)
|
||||
# x2 will be fallback allocated on the filesystem.
|
||||
x2 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
x2p = ray.get(x2)
|
||||
del x1p
|
||||
del x2p
|
||||
check_spilled_mb(address, spilled=None, fallback=400)
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_spilling_when_possible_on_put():
|
||||
try:
|
||||
address = _init_ray()
|
||||
results = []
|
||||
for _ in range(5):
|
||||
results.append(ray.put(np.zeros(400 * MB, dtype=np.uint8)))
|
||||
check_spilled_mb(address, spilled=1600)
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_fallback_when_spilling_impossible_on_get():
|
||||
try:
|
||||
address = _init_ray()
|
||||
x1 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
# x1 will be spilled.
|
||||
x2 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
check_spilled_mb(address, spilled=400)
|
||||
# x1 will be restored, x2 will be spilled.
|
||||
x1p = ray.get(x1)
|
||||
check_spilled_mb(address, spilled=800, restored=400)
|
||||
# x2 will be restored, triggering a fallback allocation.
|
||||
x2p = ray.get(x2)
|
||||
check_spilled_mb(address, spilled=800, restored=800, fallback=400)
|
||||
del x1p
|
||||
del x2p
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def fallback_allocation_mmaps():
|
||||
p = psutil.Process()
|
||||
return [
|
||||
mmap
|
||||
for mmap in p.memory_maps(grouped=False)
|
||||
if mmap.path.startswith("/tmp/ray/plasma")
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() != "Linux", reason="Using the Linux psutil.Process.memory_maps()"
|
||||
)
|
||||
def test_core_worker_fallback_allocations_munmap():
|
||||
try:
|
||||
address = _init_ray()
|
||||
x1 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
# x1 will be spilled.
|
||||
x2 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
check_spilled_mb(address, spilled=400)
|
||||
# x1 will be restored, x2 will be spilled.
|
||||
x1p = ray.get(x1)
|
||||
check_spilled_mb(address, spilled=800, restored=400)
|
||||
# No fallback allocations yet
|
||||
assert len(fallback_allocation_mmaps()) == 0, fallback_allocation_mmaps()
|
||||
# x2 will be restored, triggering a fallback allocation.
|
||||
x2p = ray.get(x2)
|
||||
check_spilled_mb(address, spilled=800, restored=800, fallback=400)
|
||||
assert len(fallback_allocation_mmaps()) == 1, fallback_allocation_mmaps()
|
||||
del x1p
|
||||
del x2p
|
||||
# after the del, the fallback allocation should be unmapped.
|
||||
assert len(fallback_allocation_mmaps()) == 0, fallback_allocation_mmaps()
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_spilling_when_possible_on_get():
|
||||
try:
|
||||
address = _init_ray()
|
||||
x1 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
# x1 will be spilled.
|
||||
x2 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
check_spilled_mb(address, spilled=400)
|
||||
# x1 will be restored, x2 will be spilled.
|
||||
ray.get(x1)
|
||||
check_spilled_mb(address, spilled=800, restored=400)
|
||||
# x2 will be restored, spilling x1.
|
||||
ray.get(x2)
|
||||
check_spilled_mb(address, spilled=800, restored=800)
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_task_unlimited():
|
||||
try:
|
||||
address = _init_ray()
|
||||
x1 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
refs = [x1]
|
||||
# x1 is spilled.
|
||||
x2 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
x2p = ray.get(x2)
|
||||
sentinel = ray.put(np.zeros(100 * MB, dtype=np.uint8))
|
||||
check_spilled_mb(address, spilled=400)
|
||||
|
||||
@ray.remote
|
||||
def consume(refs):
|
||||
# triggers fallback allocation, spilling of the sentinel
|
||||
ray.get(refs[0])
|
||||
check_spilled_mb(address, spilled=500, restored=400, fallback=400)
|
||||
# triggers fallback allocation.
|
||||
return ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
|
||||
# round 1
|
||||
_ = ray.get(ray.get(consume.remote(refs)))
|
||||
check_spilled_mb(address, spilled=500, restored=400, fallback=400)
|
||||
|
||||
del x2p
|
||||
del sentinel
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_task_unlimited_multiget_args():
|
||||
try:
|
||||
address = _init_ray()
|
||||
# Too many refs to fit into memory.
|
||||
refs = []
|
||||
for _ in range(10):
|
||||
refs.append(ray.put(np.zeros(200 * MB, dtype=np.uint8)))
|
||||
x2 = ray.put(np.zeros(600 * MB, dtype=np.uint8))
|
||||
x2p = ray.get(x2)
|
||||
check_spilled_mb(address, spilled=2000)
|
||||
|
||||
@ray.remote
|
||||
def consume(refs):
|
||||
# Should work without thrashing.
|
||||
ray.get(refs)
|
||||
return os.getpid()
|
||||
|
||||
ray.get([consume.remote(refs) for _ in range(1000)])
|
||||
check_spilled_mb(address, spilled=2000, restored=2000, fallback=2000)
|
||||
del x2p
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_fd_reuse_no_memory_corruption(shutdown_only):
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def produce(self, i):
|
||||
s = random.randrange(1, 200)
|
||||
z = np.ones(s * 1024 * 1024)
|
||||
z[0] = i
|
||||
return z
|
||||
|
||||
def consume(self, x, i):
|
||||
print(x)
|
||||
assert x[0] == i, x
|
||||
|
||||
ray.init(object_store_memory=100e6)
|
||||
a = Actor.remote()
|
||||
b = Actor.remote()
|
||||
for i in range(20):
|
||||
x_id = a.produce.remote(i)
|
||||
ray.get(b.consume.remote(x_id, i))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() != "Linux",
|
||||
reason="Only Linux handles fallback allocation disk full error.",
|
||||
)
|
||||
def test_fallback_allocation_failure(shutdown_only):
|
||||
file_system_config = {
|
||||
"type": "filesystem",
|
||||
"params": {
|
||||
"directory_path": "/dev/shm",
|
||||
},
|
||||
}
|
||||
ray.init(
|
||||
object_store_memory=100e6,
|
||||
_system_config={
|
||||
"object_spilling_config": json.dumps(file_system_config),
|
||||
# set local fs capacity to 100% so it never errors with out of disk.
|
||||
"local_fs_capacity_threshold": 1,
|
||||
},
|
||||
)
|
||||
shm_size = shutil.disk_usage("/dev/shm").total
|
||||
object_size = max(100e6, shm_size // 5)
|
||||
num_exceptions = 0
|
||||
refs = []
|
||||
for i in range(8):
|
||||
print("Start put", i)
|
||||
try:
|
||||
refs.append(ray.get(ray.put(np.zeros(object_size, dtype=np.uint8))))
|
||||
except ray.exceptions.OutOfDiskError:
|
||||
num_exceptions = num_exceptions + 1
|
||||
assert num_exceptions > 0
|
||||
|
||||
|
||||
# TODO(ekl) enable this test once we implement this behavior.
|
||||
# @pytest.mark.skipif(
|
||||
# platform.system() == "Windows", reason="Need to fix up for Windows.")
|
||||
# def test_task_unlimited_huge_args():
|
||||
# try:
|
||||
# address = _init_ray()
|
||||
#
|
||||
# # PullManager should raise an error, since the set of task args is
|
||||
# # too huge to fit into memory.
|
||||
# @ray.remote
|
||||
# def consume(*refs):
|
||||
# return "ok"
|
||||
#
|
||||
# # Too many refs to fit into memory.
|
||||
# refs = []
|
||||
# for _ in range(10):
|
||||
# refs.append(ray.put(np.zeros(200 * MB, dtype=np.uint8)))
|
||||
#
|
||||
# with pytest.raises(Exception):
|
||||
# ray.get(consume.remote(*refs))
|
||||
# finally:
|
||||
# ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_plasma_allocate(shutdown_only):
|
||||
address = ray.init(
|
||||
object_store_memory=300 * 1024**2,
|
||||
_system_config={
|
||||
"max_io_workers": 4,
|
||||
"automatic_object_spilling_enabled": True,
|
||||
},
|
||||
_temp_dir="/tmp/for_test_plasma_allocate",
|
||||
)
|
||||
res = []
|
||||
data = np.random.randint(low=0, high=256, size=(90 * 1024**2,), dtype=np.uint8)
|
||||
for _ in range(3):
|
||||
res.append(ray.put(data))
|
||||
# keep reference for second and third object, force evict first object
|
||||
_ = ray.get(res[1:]) # noqa
|
||||
# keep reference for fourth object, avoid released by plasma GC.
|
||||
__ = ray.put(data) # noqa
|
||||
|
||||
# Check fourth object allocate in memory.
|
||||
check_spilled_mb(address, spilled=[90, 180])
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
platform.system() == "Windows", reason="Need to fix up for Windows."
|
||||
)
|
||||
def test_object_store_memory_metrics_reported_correctly(shutdown_only):
|
||||
"""
|
||||
Verify when fallback allocation is used, prometheus stats report the correct
|
||||
used object store memory. https://github.com/ray-project/ray/issues/24624
|
||||
"""
|
||||
obj_store_memory = 700e6
|
||||
address = ray.init(
|
||||
num_cpus=2,
|
||||
object_store_memory=obj_store_memory,
|
||||
_system_config={"metrics_report_interval_ms": 1000},
|
||||
)
|
||||
metrics_export_port = address["metrics_export_port"]
|
||||
addr = address["node_ip_address"]
|
||||
prom_addr = build_address(addr, metrics_export_port)
|
||||
|
||||
x1 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
# x1 will be spilled.
|
||||
x2 = ray.put(np.zeros(400 * MB, dtype=np.uint8))
|
||||
check_spilled_mb(address, spilled=400)
|
||||
# x1 will be restored, x2 will be spilled.
|
||||
x1p = ray.get(x1)
|
||||
check_spilled_mb(address, spilled=800, restored=400)
|
||||
# x2 will be restored, triggering a fallback allocation.
|
||||
x2p = ray.get(x2)
|
||||
check_spilled_mb(address, spilled=800, restored=800, fallback=400)
|
||||
|
||||
def verify_used_object_store_memory(expected_mb):
|
||||
_, _, metric_samples = fetch_prometheus([prom_addr])
|
||||
|
||||
def in_mb(bytes):
|
||||
return int(bytes / 1024 / 1024)
|
||||
|
||||
total_memory = in_mb(obj_store_memory)
|
||||
available_memory_sample = None
|
||||
used_memory_sample = None
|
||||
fallback_memory_sample = None
|
||||
|
||||
for sample in metric_samples:
|
||||
if sample.name == "ray_object_store_available_memory":
|
||||
available_memory_sample = sample
|
||||
if sample.name == "ray_object_store_used_memory":
|
||||
used_memory_sample = sample
|
||||
if sample.name == "ray_object_store_fallback_memory":
|
||||
fallback_memory_sample = sample
|
||||
|
||||
if not (
|
||||
available_memory_sample and used_memory_sample and fallback_memory_sample
|
||||
):
|
||||
return False
|
||||
|
||||
avail_memory = in_mb(available_memory_sample.value)
|
||||
used_memory = in_mb(used_memory_sample.value)
|
||||
fallback_memory = in_mb(fallback_memory_sample.value)
|
||||
|
||||
assert avail_memory == total_memory - used_memory
|
||||
assert used_memory == 400 # 400MB
|
||||
assert fallback_memory == 400
|
||||
return True
|
||||
|
||||
wait_for_condition(lambda: verify_used_object_store_memory(expected_mb=30))
|
||||
del x1p
|
||||
del x2p
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
Reference in New Issue
Block a user