def test_one_item_host_limit(capsys):
memory_limit = sizeof(asproxy(one_item_array(), serializers=("dask", "pickle")))
dhf = ProxifyHostFile(
device_memory_limit=one_item_nbytes, memory_limit=memory_limit
)
a1 = one_item_array() + 1
a2 = one_item_array() + 2
dhf["k1"] = a1
dhf["k2"] = a2
dhf.manager.validate()
# Check k1 is spilled because of the newer k2
k1 = dhf["k1"]
k2 = dhf["k2"]
assert k1._pxy_get().is_serialized()
assert not k2._pxy_get().is_serialized()
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._disk.get_proxies(), [])
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k2])
# Check k1 is spilled to disk and k2 is spilled to host
dhf["k3"] = one_item_array() + 3
k3 = dhf["k3"]
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._disk.get_proxies(), [k1])
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k2])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k3])
dhf.manager.validate()
# Accessing k2 spills k3 and unspill k2
k2_val = k2[0]
assert k2_val == 2
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._disk.get_proxies(), [k1])
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k3])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k2])
# Adding a new array spill k3 to disk and k2 to host
dhf["k4"] = one_item_array() + 4
k4 = dhf["k4"]
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._disk.get_proxies(), [k1, k3])
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k2])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k4])
# Accessing k1 unspills k1 directly to device and spills k4 to host
k1_val = k1[0]
assert k1_val == 1
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._disk.get_proxies(), [k2, k3])
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k4])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k1])
# Clean up
del k1, k2, k3, k4
dhf.clear()
assert len(dhf.manager) == 0
def test_one_dev_item_limit():
dhf = ProxifyHostFile(device_memory_limit=one_item_nbytes, memory_limit=1000)
a1 = one_item_array() + 42
a2 = one_item_array()
dhf["k1"] = a1
dhf["k2"] = a2
dhf.manager.validate()
# Check k1 is spilled because of the newer k2
k1 = dhf["k1"]
k2 = dhf["k2"]
assert k1._pxy_get().is_serialized()
assert not k2._pxy_get().is_serialized()
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k2])
# Accessing k1 spills k2 and unspill k1
k1_val = k1[0]
assert k1_val == 42
assert k2._pxy_get().is_serialized()
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k2])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k1])
# Duplicate arrays changes nothing
dhf["k3"] = [k1, k2]
assert not k1._pxy_get().is_serialized()
assert k2._pxy_get().is_serialized()
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k2])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k1])
# Adding a new array spills k1 and k2
dhf["k4"] = one_item_array()
k4 = dhf["k4"]
assert k1._pxy_get().is_serialized()
assert k2._pxy_get().is_serialized()
assert not dhf["k4"]._pxy_get().is_serialized()
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1, k2])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k4])
# Accessing k2 spills k1 and k4
k2[0]
assert k1._pxy_get().is_serialized()
assert dhf["k4"]._pxy_get().is_serialized()
assert not k2._pxy_get().is_serialized()
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1, k4])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k2])
# Deleting k2 does not change anything since k3 still holds a
# reference to the underlying proxy object
assert dhf.manager._dev.mem_usage() == one_item_nbytes
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1, k4])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k2])
del dhf["k2"]
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1, k4])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k2])
# Overwriting k3 with a non-cuda object and deleting k2
# should empty the device
dhf["k3"] = "non-cuda-object"
del k2
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k1, k4])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [])
# Adding the underlying proxied of k1 doesn't change anything.
# The host file detects that k1_ary is already proxied by the
# existing proxy object k1.
k1_ary = unproxy(k1)
dhf["k5"] = k1_ary
dhf.manager.validate()
assert is_proxies_equal(dhf.manager._host.get_proxies(), [k4])
assert is_proxies_equal(dhf.manager._dev.get_proxies(), [k1])
# Clean up
del k1, k4
dhf.clear()
assert len(dhf.manager) == 0