chore: import upstream snapshot with attribution

python/ray/tune/tests/test_util_file_transfer.py

@@ -0,0 +1,247 @@
+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__]))