475 lines
17 KiB
Python
475 lines
17 KiB
Python
"""Tests for the sync r2 transfer boundary."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from sync_engine_test_support import * # noqa: F401,F403
|
|
|
|
def test_r2_worker_preflight_runs_before_transfer_and_reports_ip_filter_hint(
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
r2_cache_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A remote AccessDenied probe should become an actionable preparation failure."""
|
|
events: list[str] = []
|
|
|
|
class PreflightVolume:
|
|
"""Reject the read-only worker probe with a sanitized R2 diagnostic."""
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Satisfy the storage protocol without accessing a worker path."""
|
|
del remote_path
|
|
return False
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Reject asset mutation before a successful preflight."""
|
|
del local_path, remote_path
|
|
raise AssertionError("preflight must run before file transfer")
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Reject asset mutation before a successful preflight."""
|
|
del payload, remote_path
|
|
raise AssertionError("preflight must run before byte transfer")
|
|
|
|
def preflight_r2_access(
|
|
self,
|
|
request: Any,
|
|
*,
|
|
cancellation_check: Any = None,
|
|
) -> None:
|
|
"""Record the protected request and report the worker policy denial."""
|
|
del cancellation_check
|
|
events.append(request.allowed_host)
|
|
raise RuntimeError(
|
|
"R2 materializer failed safely "
|
|
"(category=http_client status=403 r2_code=AccessDenied)."
|
|
)
|
|
|
|
class PreflightCache:
|
|
"""Return one controller-approved read-only worker request."""
|
|
|
|
write_back_mode = "off"
|
|
|
|
def worker_preflight_request(self) -> Any:
|
|
"""Record controller validation before returning the signed probe."""
|
|
events.append("controller")
|
|
return r2_cache_module.R2WorkerPreflightRequest(
|
|
url="https://account.r2.cloudflarestorage.com/missing?secret=1",
|
|
allowed_host="account.r2.cloudflarestorage.com",
|
|
)
|
|
|
|
engine = sync_engine_module.ModalAssetSyncEngine(
|
|
volume=PreflightVolume(),
|
|
settings=_r2_sync_settings(settings_module, tmp_path),
|
|
r2_cache=PreflightCache(),
|
|
)
|
|
statuses: list[str] = []
|
|
|
|
with pytest.raises(
|
|
r2_cache_module.R2CacheError,
|
|
match="Remove Client IP Address Filtering",
|
|
):
|
|
engine.preflight_r2_access(
|
|
status_callback=lambda message, _current, _total: statuses.append(message)
|
|
)
|
|
|
|
assert events == ["controller", "account.r2.cloudflarestorage.com"]
|
|
assert statuses == [
|
|
"Checking Cloudflare R2 access from the Vast.ai instance"
|
|
]
|
|
|
|
def test_r2_cache_hit_materializes_without_local_upload(
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
r2_cache_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""An R2 hit should transfer directly to the worker and update its local index."""
|
|
asset_path = tmp_path / "model.safetensors"
|
|
asset_path.write_bytes(b"r2-backed-model")
|
|
sha256 = hashlib.sha256(asset_path.read_bytes()).hexdigest()
|
|
materializations: list[tuple[Any, str]] = []
|
|
uploads: list[tuple[Path, str]] = []
|
|
statuses: list[str] = []
|
|
|
|
class R2Volume:
|
|
"""Record remote R2 and fallback transfer operations."""
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Report an initially empty remote volume."""
|
|
del remote_path
|
|
return False
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Record a fallback local upload."""
|
|
uploads.append((local_path, remote_path))
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Accept protocol-required byte uploads."""
|
|
del payload, remote_path
|
|
|
|
def materialize_r2_file(
|
|
self,
|
|
request: Any,
|
|
remote_path: str,
|
|
*,
|
|
cancellation_check: Any = None,
|
|
) -> None:
|
|
"""Record one direct R2 materialization."""
|
|
del cancellation_check
|
|
materializations.append((request, remote_path))
|
|
|
|
def upload_r2_file(self, plan: Any, remote_path: str) -> Any:
|
|
"""Reject write-back, which is disabled in this test."""
|
|
del plan, remote_path
|
|
raise AssertionError("write-back should be disabled")
|
|
|
|
class R2Cache:
|
|
"""Return one deterministic cache hit."""
|
|
|
|
write_back_mode = "off"
|
|
|
|
def download_request(self, digest: str, size_bytes: int) -> Any:
|
|
"""Return a signed transfer request for the expected digest."""
|
|
return r2_cache_module.R2DownloadRequest(
|
|
url="https://account.r2.cloudflarestorage.com/object?secret=1",
|
|
allowed_host="account.r2.cloudflarestorage.com",
|
|
sha256=digest,
|
|
size_bytes=size_bytes,
|
|
)
|
|
|
|
engine = sync_engine_module.ModalAssetSyncEngine(
|
|
volume=R2Volume(),
|
|
settings=_r2_sync_settings(settings_module, tmp_path),
|
|
r2_cache=R2Cache(),
|
|
)
|
|
|
|
result = engine.sync_file(
|
|
asset_path,
|
|
status_callback=lambda message, current, total: statuses.append(message),
|
|
)
|
|
|
|
assert result.sha256 == sha256
|
|
assert result.uploaded is True
|
|
assert len(materializations) == 1
|
|
assert uploads == []
|
|
assert statuses == ["Downloading asset from Cloudflare R2: model.safetensors"]
|
|
|
|
def test_local_upload_queues_r2_writeback_without_blocking(
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
r2_cache_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A cache miss should return before background R2 population completes."""
|
|
asset_path = tmp_path / "model.safetensors"
|
|
asset_path.write_bytes(b"new-model")
|
|
sha256 = hashlib.sha256(asset_path.read_bytes()).hexdigest()
|
|
remote_uploads: list[str] = []
|
|
completed_plans: list[Any] = []
|
|
writeback_started = threading.Event()
|
|
allow_writeback = threading.Event()
|
|
|
|
class R2Volume:
|
|
"""Record local worker upload and R2 write-back operations."""
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Report empty worker storage."""
|
|
del remote_path
|
|
return False
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Record the normal local-to-worker upload."""
|
|
del local_path
|
|
remote_uploads.append(remote_path)
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Accept protocol-required byte uploads."""
|
|
del payload, remote_path
|
|
|
|
def materialize_r2_file(self, *args: Any, **kwargs: Any) -> None:
|
|
"""Reject a materialization on the forced cache miss."""
|
|
del args, kwargs
|
|
raise AssertionError("cache miss should not materialize")
|
|
|
|
def upload_r2_file(self, plan: Any, remote_path: str) -> Any:
|
|
"""Return a successful single-part upload result."""
|
|
assert remote_path == remote_uploads[0]
|
|
writeback_started.set()
|
|
assert allow_writeback.wait(timeout=2.0)
|
|
return r2_cache_module.R2UploadResult()
|
|
|
|
class R2Cache:
|
|
"""Expose one legacy sync setting that must migrate to background work."""
|
|
|
|
write_back_mode = "sync"
|
|
|
|
def download_request(self, digest: str, size_bytes: int) -> None:
|
|
"""Report an R2 miss."""
|
|
del digest, size_bytes
|
|
return None
|
|
|
|
def prepare_upload(
|
|
self,
|
|
digest: str,
|
|
size_bytes: int,
|
|
*,
|
|
force: bool = False,
|
|
) -> Any:
|
|
"""Return one single-part signed upload plan."""
|
|
assert force is False
|
|
return r2_cache_module.R2UploadPlan(
|
|
key=f"cache/{digest}",
|
|
sha256=digest,
|
|
size_bytes=size_bytes,
|
|
allowed_host="account.r2.cloudflarestorage.com",
|
|
mode="single",
|
|
urls=("https://account.r2.cloudflarestorage.com/object?secret=1",),
|
|
)
|
|
|
|
def complete_upload(self, plan: Any, result: Any) -> None:
|
|
"""Record successful controller completion."""
|
|
assert result == r2_cache_module.R2UploadResult()
|
|
completed_plans.append(plan)
|
|
|
|
def abort_upload(self, plan: Any) -> None:
|
|
"""Reject an abort for a successful write-back."""
|
|
del plan
|
|
raise AssertionError("successful write-back should not abort")
|
|
|
|
engine = sync_engine_module.ModalAssetSyncEngine(
|
|
volume=R2Volume(),
|
|
settings=_r2_sync_settings(settings_module, tmp_path),
|
|
r2_cache=R2Cache(),
|
|
)
|
|
|
|
result = engine.sync_file(asset_path)
|
|
assert writeback_started.wait(timeout=2.0)
|
|
assert completed_plans == []
|
|
allow_writeback.set()
|
|
engine.wait_for_r2_writebacks()
|
|
|
|
assert result.sha256 == sha256
|
|
assert len(remote_uploads) == 1
|
|
assert len(completed_plans) == 1
|
|
assert completed_plans[0].sha256 == sha256
|
|
|
|
def test_r2_writeback_coordinator_waits_for_foreground_prompt(
|
|
sync_engine_module: Any,
|
|
) -> None:
|
|
"""Idle jobs must not start while a remote prompt holds foreground priority."""
|
|
coordinator = sync_engine_module.R2WriteBackCoordinator(max_workers=1)
|
|
started = threading.Event()
|
|
completed = threading.Event()
|
|
coordinator.begin_prompt("prompt-1")
|
|
|
|
def run_writeback(cancellation_check: Any) -> None:
|
|
"""Record one coordinator-controlled cache operation."""
|
|
assert cancellation_check() is False
|
|
started.set()
|
|
completed.set()
|
|
|
|
assert coordinator.submit(("cache", "asset"), run_writeback) is True
|
|
assert started.wait(timeout=0.05) is False
|
|
|
|
coordinator.finish_prompt("prompt-1")
|
|
|
|
assert completed.wait(timeout=2.0)
|
|
assert coordinator.wait_for_idle(timeout_seconds=2.0) is True
|
|
|
|
def test_r2_writeback_coordinator_preempts_and_requeues_for_new_prompt(
|
|
sync_engine_module: Any,
|
|
) -> None:
|
|
"""A newly preparing prompt should cancel and later resume active cache work."""
|
|
coordinator = sync_engine_module.R2WriteBackCoordinator(max_workers=1)
|
|
first_started = threading.Event()
|
|
first_cancelled = threading.Event()
|
|
completed = threading.Event()
|
|
attempts = 0
|
|
|
|
def run_writeback(cancellation_check: Any) -> None:
|
|
"""Yield the first attempt and complete the retried idle attempt."""
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts == 1:
|
|
first_started.set()
|
|
assert first_cancelled.wait(timeout=2.0)
|
|
assert cancellation_check() is True
|
|
raise sync_engine_module.R2WriteBackCancelled("foreground resumed")
|
|
assert cancellation_check() is False
|
|
completed.set()
|
|
|
|
assert coordinator.submit(("cache", "asset"), run_writeback) is True
|
|
assert first_started.wait(timeout=2.0)
|
|
coordinator.begin_prompt("prompt-2")
|
|
first_cancelled.set()
|
|
assert completed.wait(timeout=0.05) is False
|
|
|
|
coordinator.finish_prompt("prompt-2")
|
|
|
|
assert completed.wait(timeout=2.0)
|
|
assert attempts == 2
|
|
assert coordinator.wait_for_idle(timeout_seconds=2.0) is True
|
|
|
|
def test_indexed_remote_payload_is_backfilled_to_r2(
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
r2_cache_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""An indexed worker file should still populate a newly enabled R2 cache."""
|
|
asset_path = tmp_path / "model.safetensors"
|
|
asset_path.write_bytes(b"previously-synced-model")
|
|
sha256 = hashlib.sha256(asset_path.read_bytes()).hexdigest()
|
|
volume = _R2BackfillVolume(r2_cache_module)
|
|
cache = _SynchronousBackfillCache(r2_cache_module)
|
|
engine = sync_engine_module.ModalAssetSyncEngine(
|
|
volume=volume,
|
|
settings=_r2_sync_settings(settings_module, tmp_path),
|
|
r2_cache=cache,
|
|
)
|
|
indexed_remote_path = "/assets/indexed-model.safetensors"
|
|
engine.sync_index.put(
|
|
engine._asset_sync_index_key(sha256),
|
|
{"remote_path": indexed_remote_path, "source": "previous sync"},
|
|
)
|
|
|
|
result = engine.sync_file(asset_path)
|
|
engine.wait_for_r2_writebacks()
|
|
|
|
assert result.uploaded is False
|
|
assert result.remote_path == indexed_remote_path
|
|
assert volume.worker_puts == []
|
|
assert [(plan.sha256, path) for plan, path in volume.r2_uploads] == [
|
|
(sha256, indexed_remote_path)
|
|
]
|
|
assert cache.requests == [(sha256, asset_path.stat().st_size, False)]
|
|
|
|
def test_existing_remote_content_is_adopted_without_reupload(
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A lost local index should be repaired from a persistent remote volume."""
|
|
asset_path = tmp_path / "model.safetensors"
|
|
asset_path.write_bytes(b"already-remote")
|
|
|
|
class ExistingVolume:
|
|
"""Expose one already populated content-addressed path."""
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Report the expected remote file as present."""
|
|
del remote_path
|
|
return True
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Reject an unnecessary duplicate upload."""
|
|
del local_path, remote_path
|
|
raise AssertionError("existing remote file should be adopted")
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Accept protocol-required byte uploads."""
|
|
del payload, remote_path
|
|
|
|
engine = sync_engine_module.ModalAssetSyncEngine(
|
|
volume=ExistingVolume(),
|
|
settings=_r2_sync_settings(settings_module, tmp_path),
|
|
)
|
|
|
|
first = engine.sync_file(asset_path)
|
|
second = engine.sync_file(asset_path)
|
|
|
|
assert first.uploaded is False
|
|
assert second.uploaded is False
|
|
assert first.remote_path == second.remote_path
|
|
|
|
def test_failed_r2_hit_is_replaced_after_fallback_upload(
|
|
settings_module: Any,
|
|
sync_engine_module: Any,
|
|
r2_cache_module: Any,
|
|
tmp_path: Path,
|
|
) -> None:
|
|
"""A same-size corrupt cache hit should trigger forced write-back after fallback."""
|
|
asset_path = tmp_path / "model.safetensors"
|
|
asset_path.write_bytes(b"healthy-local-model")
|
|
forced_values: list[bool] = []
|
|
|
|
class CorruptR2Volume:
|
|
"""Fail cached materialization, then accept normal upload and replacement."""
|
|
|
|
def exists(self, remote_path: str) -> bool:
|
|
"""Report empty worker storage."""
|
|
del remote_path
|
|
return False
|
|
|
|
def put_file(self, local_path: Path, remote_path: str) -> None:
|
|
"""Accept the established fallback transfer."""
|
|
del local_path, remote_path
|
|
|
|
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
|
"""Accept protocol-required byte uploads."""
|
|
del payload, remote_path
|
|
|
|
def materialize_r2_file(self, *args: Any, **kwargs: Any) -> None:
|
|
"""Emulate a worker-side SHA mismatch."""
|
|
del args, kwargs
|
|
raise ValueError("R2 download digest mismatch")
|
|
|
|
def upload_r2_file(self, plan: Any, remote_path: str) -> Any:
|
|
"""Accept the forced cache replacement."""
|
|
del plan, remote_path
|
|
return r2_cache_module.R2UploadResult()
|
|
|
|
class CorruptR2Cache:
|
|
"""Return a hit and record whether replacement bypasses size reuse."""
|
|
|
|
write_back_mode = "sync"
|
|
|
|
def download_request(self, digest: str, size_bytes: int) -> Any:
|
|
"""Return one apparently valid same-size cache object."""
|
|
return r2_cache_module.R2DownloadRequest(
|
|
url="https://account.r2.cloudflarestorage.com/object?secret=1",
|
|
allowed_host="account.r2.cloudflarestorage.com",
|
|
sha256=digest,
|
|
size_bytes=size_bytes,
|
|
)
|
|
|
|
def prepare_upload(
|
|
self,
|
|
digest: str,
|
|
size_bytes: int,
|
|
*,
|
|
force: bool = False,
|
|
) -> Any:
|
|
"""Record and return one replacement plan."""
|
|
forced_values.append(force)
|
|
return r2_cache_module.R2UploadPlan(
|
|
key=f"cache/{digest}",
|
|
sha256=digest,
|
|
size_bytes=size_bytes,
|
|
allowed_host="account.r2.cloudflarestorage.com",
|
|
mode="single",
|
|
urls=("https://account.r2.cloudflarestorage.com/object?secret=2",),
|
|
)
|
|
|
|
def complete_upload(self, plan: Any, result: Any) -> None:
|
|
"""Accept successful forced replacement completion."""
|
|
del plan, result
|
|
|
|
def abort_upload(self, plan: Any) -> None:
|
|
"""Reject an abort in the successful path."""
|
|
del plan
|
|
raise AssertionError("replacement should not abort")
|
|
|
|
engine = sync_engine_module.ModalAssetSyncEngine(
|
|
volume=CorruptR2Volume(),
|
|
settings=_r2_sync_settings(settings_module, tmp_path),
|
|
r2_cache=CorruptR2Cache(),
|
|
)
|
|
|
|
engine.sync_file(asset_path)
|
|
engine.wait_for_r2_writebacks()
|
|
|
|
assert forced_values == [True]
|
|
|