def test_device_host_file_short(tmp_path, num_device_arrays, num_host_arrays, array_size_range): tmpdir = tmp_path / "storage" tmpdir.mkdir() dhf = DeviceHostFile(device_memory_limit=1024 * 16, memory_limit=1024 * 16, local_dir=tmpdir) host = [("x-%d" % i, np.random.random(randint(*array_size_range))) for i in range(num_host_arrays)] device = [("dx-%d" % i, cupy.random.random(randint(*array_size_range))) for i in range(num_device_arrays)] import random full = host + device random.shuffle(full) for i in full: dhf[i[0]] = i[1] random.shuffle(full) for i in full: assert_array_equal(i[1], dhf[i[0]]) del dhf[i[0]] assert set(dhf.device.keys()) == set() assert set(dhf.host.keys()) == set() assert set(dhf.disk.keys()) == set()
def test_device_host_file_short(tmp_path, num_device_arrays, num_host_arrays, array_size_range, jit_unspill): tmpdir = tmp_path / "storage" tmpdir.mkdir() dhf = DeviceHostFile( device_memory_limit=1024 * 16, memory_limit=1024 * 16, local_directory=tmpdir, jit_unspill=jit_unspill, ) host = [("x-%d" % i, np.random.random(randint(*array_size_range))) for i in range(num_host_arrays)] device = [("dx-%d" % i, cupy.random.random(randint(*array_size_range))) for i in range(num_device_arrays)] import random full = host + device random.shuffle(full) for k, v in full: dhf[k] = v random.shuffle(full) for k, original in full: acquired = dhf[k] assert_eq(original, acquired) del dhf[k] assert set(dhf.device.keys()) == set() assert set(dhf.host.keys()) == set() assert set(dhf.disk.keys()) == set()
def test_cudf_device_spill(params): @gen_cluster( client=True, ncores=[("127.0.0.1", 1)], Worker=Worker, timeout=300, worker_kwargs={ "memory_limit": params["memory_limit"], "data": DeviceHostFile( device_memory_limit=params["device_memory_limit"], memory_limit=params["memory_limit"], ), }, config={ "distributed.comm.timeouts.connect": "20s", "distributed.worker.memory.target": params["host_target"], "distributed.worker.memory.spill": params["host_spill"], }, ) def test_device_spill(client, scheduler, worker): cdf = dask.datasets.timeseries(dtypes={ "x": int, "y": float }, freq="30ms").map_partitions( cudf.from_pandas) sizes = yield client.compute( cdf.map_partitions(lambda df: df.__sizeof__())) sizes = sizes.tolist() nbytes = sum(sizes) part_index_nbytes = (yield client.compute( cdf.partitions[0].index)).__sizeof__() cdf2 = cdf.persist() yield wait(cdf2) del cdf yield client.run(worker_assert, nbytes, 32, 2048 + part_index_nbytes) host_chunks = yield client.run(lambda: len(get_worker().data.host)) disk_chunks = yield client.run(lambda: len(get_worker().data.disk)) for hc, dc in zip(host_chunks.values(), disk_chunks.values()): if params["spills_to_disk"]: assert dc > 0 else: assert hc > 0 assert dc == 0 del cdf2 yield client.run(worker_assert, 0, 0, 0) test_device_spill()
def test_cudf_device_spill(params): @gen_cluster( client=True, nthreads=[("127.0.0.1", 1)], Worker=Worker, timeout=60, worker_kwargs={ "memory_limit": params["memory_limit"], "data": DeviceHostFile( device_memory_limit=params["device_memory_limit"], memory_limit=params["memory_limit"], ), }, config={ "distributed.comm.timeouts.connect": "20s", "distributed.worker.memory.target": params["host_target"], "distributed.worker.memory.spill": params["host_spill"], "distributed.worker.memory.pause": params["host_pause"], }, ) def test_device_spill(client, scheduler, worker): cudf = pytest.importorskip("cudf") # There's a known issue with datetime64: # https://github.com/numpy/numpy/issues/4983#issuecomment-441332940 # The same error above happens when spilling datetime64 to disk cdf = ( dask.datasets.timeseries(dtypes={"x": int, "y": float}, freq="20ms") .reset_index(drop=True) .map_partitions(cudf.from_pandas) ) sizes = yield client.compute(cdf.map_partitions(lambda df: df.__sizeof__())) sizes = sizes.tolist() nbytes = sum(sizes) part_index_nbytes = (yield client.compute(cdf.partitions[0].index)).__sizeof__() cdf2 = cdf.persist() yield wait(cdf2) del cdf host_chunks = yield client.run(lambda: len(get_worker().data.host)) disk_chunks = yield client.run(lambda: len(get_worker().data.disk or list())) for hc, dc in zip(host_chunks.values(), disk_chunks.values()): if params["spills_to_disk"]: assert dc > 0 else: assert hc > 0 assert dc == 0 yield client.run(worker_assert, nbytes, 32, 2048 + part_index_nbytes) del cdf2 yield client.run(delayed_worker_assert, 0, 0, 0) test_device_spill()
def test_cupy_device_spill(params): @gen_cluster( client=True, nthreads=[("127.0.0.1", 1)], Worker=Worker, timeout=60, worker_kwargs={ "memory_limit": params["memory_limit"], "data": DeviceHostFile( device_memory_limit=params["device_memory_limit"], memory_limit=params["memory_limit"], ), }, config={ "distributed.comm.timeouts.connect": "20s", "distributed.worker.memory.target": params["host_target"], "distributed.worker.memory.spill": params["host_spill"], "distributed.worker.memory.pause": params["host_pause"], }, ) def test_device_spill(client, scheduler, worker): cupy = pytest.importorskip("cupy") rs = da.random.RandomState(RandomState=cupy.random.RandomState) x = rs.random(int(250e6), chunks=10e6) xx = x.persist() yield wait(xx) # Allow up to 1024 bytes overhead per chunk serialized yield client.run(worker_assert, x.nbytes, 1024, 1024) y = client.compute(x.sum()) res = yield y assert (abs(res / x.size) - 0.5) < 1e-3 yield client.run(worker_assert, x.nbytes, 1024, 1024) host_chunks = yield client.run(lambda: len(get_worker().data.host)) disk_chunks = yield client.run( lambda: len(get_worker().data.disk or list())) for hc, dc in zip(host_chunks.values(), disk_chunks.values()): if params["spills_to_disk"]: assert dc > 0 else: assert hc > 0 assert dc == 0 test_device_spill()
def test_device_host_file_step_by_step(tmp_path): tmpdir = tmp_path / "storage" tmpdir.mkdir() dhf = DeviceHostFile(device_memory_limit=1024 * 16, memory_limit=1024 * 16, local_dir=tmpdir) a = np.random.random(1000) b = cupy.random.random(1000) dhf["a1"] = a assert set(dhf.device.keys()) == set() assert set(dhf.host.keys()) == set(["a1"]) assert set(dhf.disk.keys()) == set() dhf["b1"] = b assert set(dhf.device.keys()) == set(["b1"]) assert set(dhf.host.keys()) == set(["a1"]) assert set(dhf.disk.keys()) == set() dhf["b2"] = b assert set(dhf.device.keys()) == set(["b1", "b2"]) assert set(dhf.host.keys()) == set(["a1"]) assert set(dhf.disk.keys()) == set() dhf["b3"] = b assert set(dhf.device.keys()) == set(["b2", "b3"]) assert set(dhf.host.keys()) == set(["a1", "b1"]) assert set(dhf.disk.keys()) == set() dhf["a2"] = a assert set(dhf.device.keys()) == set(["b2", "b3"]) assert set(dhf.host.keys()) == set(["a2", "b1"]) assert set(dhf.disk.keys()) == set(["a1"]) dhf["b4"] = b assert set(dhf.device.keys()) == set(["b3", "b4"]) assert set(dhf.host.keys()) == set(["a2", "b2"]) assert set(dhf.disk.keys()) == set(["a1", "b1"]) dhf["b4"] = b assert set(dhf.device.keys()) == set(["b3", "b4"]) assert set(dhf.host.keys()) == set(["a2", "b2"]) assert set(dhf.disk.keys()) == set(["a1", "b1"]) assert_array_equal(dhf["a1"], a) del dhf["a1"] assert_array_equal(dhf["a2"], a) del dhf["a2"] assert_array_equal(dhf["b1"], b) del dhf["b1"] assert_array_equal(dhf["b2"], b) del dhf["b2"] assert_array_equal(dhf["b3"], b) del dhf["b3"] assert_array_equal(dhf["b4"], b) del dhf["b4"] assert set(dhf.device.keys()) == set() assert set(dhf.host.keys()) == set() assert set(dhf.disk.keys()) == set()
def test_device_host_file_config(tmp_path): dhf_disk_path = str(tmp_path / "dask-worker-space" / "storage") with dask.config.set(temporary_directory=str(tmp_path)): dhf = DeviceHostFile() assert os.path.exists(dhf_disk_path) assert dhf.disk_func_path == dhf_disk_path