wehub-resource-sync
405 lines
12 KiB
Python
405 lines
12 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
from typing import List, Optional
|
|
|
|
import pytest
|
|
from freezegun import freeze_time
|
|
|
|
import ray
|
|
import ray.cloudpickle as pickle
|
|
from ray.train._internal.storage import (
|
|
_download_from_fs_path,
|
|
_FilesystemSyncer,
|
|
_upload_to_fs_path,
|
|
get_fs_and_path,
|
|
)
|
|
from ray.train._internal.syncer import _BackgroundProcess
|
|
from ray.train.tests.test_new_persistence import _create_mock_custom_fs
|
|
|
|
|
|
@pytest.fixture
|
|
def propagate_logs():
|
|
# Ensure that logs are propagated to ancestor handles. This is required if using the
|
|
# caplog fixture with Ray's logging.
|
|
# NOTE: This only enables log propagation in the driver process, not the workers!
|
|
logger = logging.getLogger("ray")
|
|
logger.propagate = True
|
|
yield
|
|
logger.propagate = False
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_4_cpus():
|
|
address_info = ray.init(num_cpus=4, configure_logging=False)
|
|
yield address_info
|
|
# The code after the yield will run as teardown code.
|
|
ray.shutdown()
|
|
|
|
|
|
@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 shutdown_only():
|
|
yield None
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_data_dirs(tmp_path):
|
|
tmp_source = tmp_path / "source"
|
|
tmp_target = tmp_path / "target"
|
|
tmp_target.mkdir()
|
|
|
|
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 str(tmp_source), str(tmp_target)
|
|
|
|
|
|
@pytest.fixture
|
|
def syncer(tmp_path):
|
|
yield _FilesystemSyncer(storage_filesystem=_create_mock_custom_fs(tmp_path))
|
|
|
|
|
|
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_syncer_sync_up(temp_data_dirs, syncer):
|
|
"""Check that syncing up works"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
syncer.sync_up(local_dir=tmp_source, remote_dir="/test/test_syncer_sync_up_down")
|
|
syncer.wait()
|
|
|
|
_download_from_fs_path(
|
|
syncer.storage_filesystem, "/test/test_syncer_sync_up_down", tmp_target
|
|
)
|
|
|
|
# Target dir should have all files
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(True, tmp_target, "level0_exclude.txt")
|
|
assert_file(True, tmp_target, "subdir/level1.txt")
|
|
assert_file(True, tmp_target, "subdir/level1_exclude.txt")
|
|
assert_file(True, tmp_target, "subdir/nested/level2.txt")
|
|
assert_file(True, tmp_target, "subdir_nested_level2_exclude.txt")
|
|
assert_file(True, tmp_target, "subdir_exclude/something/somewhere.txt")
|
|
|
|
|
|
def test_syncer_sync_exclude(temp_data_dirs, syncer):
|
|
"""Check that the exclude parameter works"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
syncer.sync_up(
|
|
local_dir=tmp_source,
|
|
remote_dir="/test/test_syncer_sync_exclude",
|
|
exclude=["*_exclude*"],
|
|
)
|
|
syncer.wait()
|
|
|
|
_download_from_fs_path(
|
|
syncer.storage_filesystem, "/test/test_syncer_sync_exclude", tmp_target
|
|
)
|
|
|
|
# Excluded files should not be found in target
|
|
assert_file(True, tmp_target, "level0.txt")
|
|
assert_file(False, tmp_target, "level0_exclude.txt")
|
|
assert_file(True, tmp_target, "subdir/level1.txt")
|
|
assert_file(False, tmp_target, "subdir/level1_exclude.txt")
|
|
assert_file(True, tmp_target, "subdir/nested/level2.txt")
|
|
assert_file(False, tmp_target, "subdir_nested_level2_exclude.txt")
|
|
assert_file(False, tmp_target, "subdir_exclude/something/somewhere.txt")
|
|
|
|
|
|
def test_sync_up_if_needed(temp_data_dirs, tmp_path):
|
|
"""Check that we only sync up again after sync period"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
with freeze_time() as frozen:
|
|
syncer = _FilesystemSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
|
|
)
|
|
|
|
assert syncer.sync_up_if_needed(
|
|
local_dir=tmp_source, remote_dir="/test/test_sync_up_not_needed"
|
|
)
|
|
syncer.wait()
|
|
|
|
frozen.tick(30)
|
|
|
|
# Sync period not over, yet
|
|
assert not syncer.sync_up_if_needed(
|
|
local_dir=tmp_source, remote_dir="/test/test_sync_up_not_needed"
|
|
)
|
|
|
|
frozen.tick(30)
|
|
|
|
# Sync period over, sync again
|
|
assert syncer.sync_up_if_needed(
|
|
local_dir=tmp_source, remote_dir="/test/test_sync_up_not_needed"
|
|
)
|
|
|
|
|
|
def test_syncer_still_running_no_sync(temp_data_dirs, tmp_path):
|
|
"""Check that no new sync is issued if old sync is still running"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
class FakeSyncProcess:
|
|
@property
|
|
def is_running(self):
|
|
return True
|
|
|
|
@property
|
|
def start_time(self):
|
|
# Don't consider the sync process timeout
|
|
return float("inf")
|
|
|
|
syncer = _FilesystemSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
|
|
)
|
|
syncer._sync_process = FakeSyncProcess()
|
|
assert not syncer.sync_up_if_needed(
|
|
local_dir=tmp_source,
|
|
remote_dir="/test/test_syncer_still_running_no_sync",
|
|
)
|
|
|
|
|
|
def test_syncer_not_running_sync(temp_data_dirs, tmp_path):
|
|
"""Check that new sync is issued if old sync completed"""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
class FakeSyncProcess:
|
|
@property
|
|
def is_running(self):
|
|
return False
|
|
|
|
def wait(self):
|
|
return True
|
|
|
|
syncer = _FilesystemSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
|
|
)
|
|
syncer._sync_process = FakeSyncProcess()
|
|
assert syncer.sync_up_if_needed(
|
|
local_dir=tmp_source,
|
|
remote_dir="/test/test_syncer_not_running_sync",
|
|
)
|
|
|
|
|
|
def test_syncer_hanging_sync_with_timeout(temp_data_dirs, tmp_path):
|
|
"""Check that syncing times out when the sync process is hanging."""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
def _hanging_sync_up_command(*args, **kwargs):
|
|
time.sleep(200)
|
|
|
|
class _HangingSyncer(_FilesystemSyncer):
|
|
def _sync_up_command(
|
|
self, local_path: str, uri: str, exclude: Optional[List] = None
|
|
):
|
|
return _hanging_sync_up_command, {}
|
|
|
|
syncer = _HangingSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path),
|
|
sync_period=60,
|
|
sync_timeout=10,
|
|
)
|
|
|
|
def sync_up():
|
|
return syncer.sync_up(
|
|
local_dir=tmp_source, remote_dir="/test/test_syncer_timeout"
|
|
)
|
|
|
|
with freeze_time() as frozen:
|
|
assert sync_up()
|
|
frozen.tick(5)
|
|
# 5 seconds - initial sync hasn't reached the timeout yet
|
|
# It should continue running without launching a new sync
|
|
assert not sync_up()
|
|
frozen.tick(5)
|
|
# Reached the timeout - start running a new sync command
|
|
assert sync_up()
|
|
frozen.tick(20)
|
|
# We're 10 seconds past the timeout, waiting should result in a timeout error
|
|
with pytest.raises(TimeoutError):
|
|
syncer.wait()
|
|
|
|
|
|
def test_syncer_not_running_sync_last_failed(
|
|
propagate_logs, caplog, temp_data_dirs, tmp_path
|
|
):
|
|
"""Check that new sync is issued if old sync completed"""
|
|
caplog.set_level(logging.WARNING)
|
|
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
class FakeSyncProcess(_BackgroundProcess):
|
|
@property
|
|
def is_running(self):
|
|
return False
|
|
|
|
def wait(self, *args, **kwargs):
|
|
raise RuntimeError("Sync failed")
|
|
|
|
syncer = _FilesystemSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path), sync_period=60
|
|
)
|
|
syncer._sync_process = FakeSyncProcess(lambda: None)
|
|
assert syncer.sync_up_if_needed(
|
|
local_dir=tmp_source,
|
|
remote_dir="/test/test_syncer_not_running_sync",
|
|
)
|
|
assert "Last sync command failed" in caplog.text
|
|
|
|
|
|
def test_syncer_wait_or_retry_failure(temp_data_dirs, tmp_path):
|
|
"""Check that the wait or retry API fails after max_retries."""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
syncer = _FilesystemSyncer(storage_filesystem=lambda: "error", sync_period=60)
|
|
|
|
# Will fail since the storage filesystem is invalid
|
|
syncer.sync_up(local_dir=tmp_source, remote_dir="/test/test_syncer_wait_or_retry")
|
|
with pytest.raises(RuntimeError) as e:
|
|
syncer.wait_or_retry(max_retries=3, backoff_s=0)
|
|
|
|
assert "Failed sync even after 3 retries." in str(e.value)
|
|
|
|
|
|
def test_syncer_wait_or_retry_timeout(temp_data_dirs, tmp_path):
|
|
"""Check that the wait or retry API raises a timeout error after `sync_timeout`."""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
def slow_upload(*args, **kwargs):
|
|
time.sleep(5)
|
|
|
|
class HangingSyncer(_FilesystemSyncer):
|
|
def _sync_up_command(
|
|
self, local_path: str, uri: str, exclude: Optional[List] = None
|
|
):
|
|
return (
|
|
slow_upload,
|
|
dict(local_path=local_path, uri=uri, exclude=exclude),
|
|
)
|
|
|
|
syncer = HangingSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path),
|
|
sync_period=60,
|
|
sync_timeout=0.1,
|
|
)
|
|
|
|
syncer.sync_up(local_dir=tmp_source, remote_dir="/test/timeout")
|
|
with pytest.raises(RuntimeError) as e:
|
|
syncer.wait_or_retry(max_retries=3, backoff_s=0)
|
|
assert "Failed sync even after 3 retries." in str(e.value)
|
|
assert isinstance(e.value.__cause__, TimeoutError)
|
|
|
|
|
|
def test_syncer_wait_or_retry_eventual_success(temp_data_dirs, tmp_path):
|
|
"""Check that the wait or retry API succeeds for a sync_down that
|
|
fails, times out, then succeeds."""
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
success = tmp_path / "success"
|
|
fail_marker = tmp_path / "fail_marker"
|
|
hang_marker = tmp_path / "hang_marker"
|
|
|
|
def eventual_upload(*args, **kwargs):
|
|
if not fail_marker.exists():
|
|
fail_marker.write_text(".", encoding="utf-8")
|
|
raise RuntimeError("Failing")
|
|
elif not hang_marker.exists():
|
|
hang_marker.write_text(".", encoding="utf-8")
|
|
time.sleep(5)
|
|
else:
|
|
success.write_text(".", encoding="utf-8")
|
|
|
|
class EventualSuccessSyncer(_FilesystemSyncer):
|
|
def _sync_up_command(
|
|
self, local_path: str, uri: str, exclude: Optional[List] = None
|
|
):
|
|
return (
|
|
eventual_upload,
|
|
dict(local_path=local_path, uri=uri, exclude=exclude),
|
|
)
|
|
|
|
syncer = EventualSuccessSyncer(
|
|
storage_filesystem=_create_mock_custom_fs(tmp_path),
|
|
sync_period=60,
|
|
sync_timeout=0.5,
|
|
)
|
|
|
|
syncer.sync_up(local_dir=tmp_source, remote_dir="/test/eventual_success")
|
|
# The syncer will retry 2 times, running 3 times in total and eventually succeeding.
|
|
syncer.wait_or_retry(max_retries=2, backoff_s=0)
|
|
assert success.exists()
|
|
|
|
|
|
def test_syncer_serialize(temp_data_dirs, syncer):
|
|
tmp_source, tmp_target = temp_data_dirs
|
|
|
|
syncer.sync_up(local_dir=tmp_source, remote_dir="/test/serialize")
|
|
|
|
serialized = pickle.dumps(syncer)
|
|
loaded_syncer = pickle.loads(serialized)
|
|
assert not loaded_syncer._sync_process
|
|
|
|
|
|
def test_sync_many_files_local_to_cloud(mock_s3_bucket_uri, tmp_path):
|
|
source_dir = tmp_path / "source"
|
|
check_dir = tmp_path / "check"
|
|
source_dir.mkdir()
|
|
check_dir.mkdir()
|
|
|
|
# Create 256 files to upload
|
|
for i in range(256):
|
|
(source_dir / str(i)).write_text("", encoding="utf-8")
|
|
|
|
fs, fs_path = get_fs_and_path(mock_s3_bucket_uri)
|
|
_upload_to_fs_path(source_dir, fs, fs_path)
|
|
_download_from_fs_path(fs, fs_path, check_dir)
|
|
assert (check_dir / "255").exists()
|
|
|
|
|
|
def test_sync_many_files_local_to_local(tmp_path):
|
|
(tmp_path / "source").mkdir()
|
|
# Create 256 files to upload
|
|
for i in range(256):
|
|
(tmp_path / "source" / str(i)).write_text("", encoding="utf-8")
|
|
|
|
fs, fs_path = get_fs_and_path(str(tmp_path / "destination"))
|
|
_upload_to_fs_path(str(tmp_path / "source"), fs, fs_path)
|
|
assert (tmp_path / "destination" / "255").exists()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|