mirror of
https://github.com/wassname/ray.git
synced 2026-06-29 09:06:10 +08:00
Zhuohan Li
8a76f4cbb5
* remove the put in memory store * put small objects directly in memory store * cast data type * fix another place that uses Put to spill to plasma store * fix multiple tests related to memory limits * partially fix test_metrics * remove not functioning codes * fix core_worker_test * refactor put to plasma codes * add a flag for the new feature * add flag to more places * do a warmup round for the plasma store * lint * lint again * fix warmup store * Update _raylet.pyx Co-authored-by: Eric Liang <ekhliang@gmail.com>
89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
import numpy as np
|
|
import unittest
|
|
|
|
import ray
|
|
|
|
MB = 1024 * 1024
|
|
|
|
OBJECT_EVICTED = ray.exceptions.UnreconstructableError
|
|
OBJECT_TOO_LARGE = ray.exceptions.ObjectStoreFullError
|
|
|
|
|
|
@ray.remote
|
|
class LightActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def sample(self):
|
|
return np.zeros(5 * MB, dtype=np.uint8)
|
|
|
|
|
|
@ray.remote
|
|
class GreedyActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def sample(self):
|
|
return np.zeros(20 * MB, dtype=np.uint8)
|
|
|
|
|
|
class TestMemoryLimits(unittest.TestCase):
|
|
def testWithoutQuota(self):
|
|
self._run(100 * MB, None, None)
|
|
self.assertRaises(OBJECT_EVICTED, lambda: self._run(None, None, None))
|
|
self.assertRaises(OBJECT_EVICTED,
|
|
lambda: self._run(None, 100 * MB, None))
|
|
|
|
def testQuotasProtectSelf(self):
|
|
self._run(100 * MB, 100 * MB, None)
|
|
|
|
def testQuotasProtectOthers(self):
|
|
self._run(None, None, 100 * MB)
|
|
|
|
def testQuotaTooLarge(self):
|
|
self.assertRaisesRegexp(ray.memory_monitor.RayOutOfMemoryError,
|
|
".*Failed to set object_store_memory.*",
|
|
lambda: self._run(300 * MB, None, None))
|
|
|
|
def testTooLargeAllocation(self):
|
|
try:
|
|
ray.init(num_cpus=1, driver_object_store_memory=100 * MB)
|
|
ray.put(np.zeros(50 * MB, dtype=np.uint8), weakref=True)
|
|
self.assertRaises(
|
|
OBJECT_TOO_LARGE,
|
|
lambda: ray.put(np.zeros(200 * MB, dtype=np.uint8)))
|
|
finally:
|
|
ray.shutdown()
|
|
|
|
def _run(self, driver_quota, a_quota, b_quota):
|
|
print("*** Testing ***", driver_quota, a_quota, b_quota)
|
|
try:
|
|
ray.init(
|
|
num_cpus=1,
|
|
object_store_memory=300 * MB,
|
|
driver_object_store_memory=driver_quota)
|
|
obj = np.ones(200 * 1024, dtype=np.uint8)
|
|
z = ray.put(obj, weakref=True)
|
|
a = LightActor._remote(object_store_memory=a_quota)
|
|
b = GreedyActor._remote(object_store_memory=b_quota)
|
|
for _ in range(5):
|
|
r_a = a.sample.remote()
|
|
for _ in range(20):
|
|
new_oid = b.sample.remote()
|
|
ray.get(new_oid)
|
|
ray.get(r_a)
|
|
ray.get(z)
|
|
except Exception as e:
|
|
print("Raised exception", type(e), e)
|
|
raise e
|
|
finally:
|
|
print(ray.worker.global_worker.core_worker.
|
|
dump_object_store_memory_usage())
|
|
ray.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
import sys
|
|
sys.exit(pytest.main(["-v", __file__]))
|