wehub-resource-sync
248 lines
7.5 KiB
Python
248 lines
7.5 KiB
Python
import io
|
|
import os
|
|
import shutil
|
|
import tarfile
|
|
import tempfile
|
|
|
|
import pytest
|
|
|
|
import ray.util
|
|
from ray.exceptions import RayTaskError
|
|
from ray.tune.utils.file_transfer import (
|
|
_sync_dir_between_different_nodes,
|
|
_sync_dir_on_same_node,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_2_cpus():
|
|
address_info = ray.init(num_cpus=2, configure_logging=False)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_data_dirs():
|
|
tmp_source = os.path.realpath(tempfile.mkdtemp())
|
|
tmp_target = os.path.realpath(tempfile.mkdtemp())
|
|
|
|
os.makedirs(os.path.join(tmp_source, "subdir", "nested"))
|
|
os.makedirs(os.path.join(tmp_source, "subdir_exclude", "something"))
|
|
|
|
files = [
|
|
"level0.txt",
|
|
"level0_exclude.txt",
|
|
"subdir/level1.txt",
|
|
"subdir/level1_exclude.txt",
|
|
"subdir/nested/level2.txt",
|
|
"subdir_nested_level2_exclude.txt",
|
|
"subdir_exclude/something/somewhere.txt",
|
|
]
|
|
|
|
for file in files:
|
|
with open(os.path.join(tmp_source, file), "w") as f:
|
|
f.write("Data")
|
|
|
|
yield tmp_source, tmp_target
|
|
|
|
shutil.rmtree(tmp_source)
|
|
shutil.rmtree(tmp_target)
|
|
|
|
|
|
def assert_file(exists: bool, root: str, path: str):
|
|
full_path = os.path.join(root, path)
|
|
|
|
if exists:
|
|
assert os.path.exists(full_path)
|
|
else:
|
|
assert not os.path.exists(full_path)
|
|
|
|
|
|
def test_sync_nodes(ray_start_2_cpus, temp_data_dirs):
|
|
"""Check that syncing between nodes works (data is found in target directory)"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_source, "subdir/level1.txt")
|
|
assert_file(False, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "subdir/level1.txt")
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
_sync_dir_between_different_nodes(
|
|
source_ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_ip=node_ip,
|
|
target_path=tmp_target,
|
|
)
|
|
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(True, tmp_target, "subdir/level1.txt")
|
|
|
|
|
|
def test_sync_nodes_only_diff(ray_start_2_cpus, temp_data_dirs):
|
|
"""Check that only differing files are synced between nodes"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
# Sanity check
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_source, "subdir/level1.txt")
|
|
assert_file(False, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "level0_new.txt")
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
_sync_dir_between_different_nodes(
|
|
source_ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_ip=node_ip,
|
|
target_path=tmp_target,
|
|
)
|
|
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(True, tmp_target, "subdir/level1.txt")
|
|
assert_file(False, tmp_target, "level0_new.txt")
|
|
|
|
# Add new file
|
|
with open(os.path.join(tmp_source, "level0_new.txt"), "w") as f:
|
|
f.write("Data\n")
|
|
|
|
# Modify existing file
|
|
with open(os.path.join(tmp_source, "subdir", "level1.txt"), "w") as f:
|
|
f.write("New data\n")
|
|
|
|
unpack, pack_actor, files_stats = _sync_dir_between_different_nodes(
|
|
source_ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_ip=node_ip,
|
|
target_path=tmp_target,
|
|
return_futures=True,
|
|
)
|
|
|
|
files_stats = ray.get(files_stats)
|
|
tarball = ray.get(pack_actor.get_full_data.remote())
|
|
|
|
assert "./level0.txt" in files_stats
|
|
assert "./level0_new.txt" not in files_stats # Was not in target dir
|
|
assert "subdir/level1.txt" in files_stats
|
|
|
|
with tarfile.open(fileobj=io.BytesIO(tarball)) as tar:
|
|
files_in_tar = tar.getnames()
|
|
assert "./level0.txt" not in files_in_tar
|
|
assert "./level0_new.txt" in files_in_tar
|
|
assert "subdir/level1.txt" in files_in_tar
|
|
assert len(files_in_tar) == 7 # 3 files, 4 dirs (including root)
|
|
|
|
ray.get(unpack) # Wait until finished for teardown
|
|
|
|
|
|
@pytest.mark.parametrize("exclude", [["subdir/*"], ["*/level1.txt"]])
|
|
def test_sync_nodes_exclude_different_node(ray_start_2_cpus, temp_data_dirs, exclude):
|
|
"""Check that excluding files works"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_source, "subdir/level1.txt")
|
|
assert_file(False, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "subdir/level1.txt")
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
_sync_dir_between_different_nodes(
|
|
source_ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_ip=node_ip,
|
|
target_path=tmp_target,
|
|
exclude=exclude,
|
|
)
|
|
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "subdir/level1.txt")
|
|
|
|
|
|
@pytest.mark.parametrize("exclude", [["subdir/*"], ["*/level1.txt"]])
|
|
def test_sync_nodes_exclude_same_node(ray_start_2_cpus, temp_data_dirs, exclude):
|
|
"""Check that excluding files works"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_source, "subdir/level1.txt")
|
|
assert_file(False, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "subdir/level1.txt")
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
_sync_dir_on_same_node(
|
|
ip=node_ip, source_path=tmp_source, target_path=tmp_target, exclude=exclude
|
|
)
|
|
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "subdir/level1.txt")
|
|
|
|
|
|
@pytest.mark.parametrize("num_workers", [1, 8])
|
|
def test_multi_sync_same_node(ray_start_2_cpus, temp_data_dirs, num_workers):
|
|
"""Check that multiple competing syncs to the same node+dir don't interfere"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_source, "subdir/level1.txt")
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
futures = [
|
|
_sync_dir_on_same_node(
|
|
ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_path=tmp_target,
|
|
return_futures=True,
|
|
)
|
|
for _ in range(num_workers)
|
|
]
|
|
ray.get(futures)
|
|
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(True, tmp_target, "subdir/level1.txt")
|
|
|
|
|
|
@pytest.mark.parametrize("num_workers", [1, 8])
|
|
def test_multi_sync_different_node(ray_start_2_cpus, temp_data_dirs, num_workers):
|
|
"""Check that multiple competing syncs to the same dir don't interfere"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
assert_file(True, tmp_source, "level0.txt")
|
|
assert_file(True, tmp_source, "subdir/level1.txt")
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
futures = [
|
|
_sync_dir_between_different_nodes(
|
|
source_ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_ip=node_ip,
|
|
target_path=tmp_target,
|
|
return_futures=True,
|
|
)[0]
|
|
for _ in range(num_workers)
|
|
]
|
|
ray.get(futures)
|
|
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(True, tmp_target, "subdir/level1.txt")
|
|
|
|
|
|
def test_max_size_exceeded(ray_start_2_cpus, temp_data_dirs):
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
node_ip = ray.util.get_node_ip_address()
|
|
with pytest.raises(RayTaskError):
|
|
_sync_dir_between_different_nodes(
|
|
source_ip=node_ip,
|
|
source_path=tmp_source,
|
|
target_ip=node_ip,
|
|
target_path=tmp_target,
|
|
max_size_bytes=2,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|