def run_storage_lifecycle_test(
manager: storage.StorageManager,
post_delete_cb: Optional[Callable] = None,
) -> None:
checkpoints = []
for _ in range(2):
storage_id = str(uuid.uuid4())
with manager.store_path(storage_id) as path:
create_checkpoint(path)
checkpoints.append(storage_id)
for storage_id in checkpoints:
# Load checkpoint.
with manager.restore_path(storage_id) as path:
validate_checkpoint(path)
# Delete.
manager.delete(storage_id)
# Ensure it is gone.
with pytest.raises(errors.CheckpointNotFound):
with manager.restore_path(storage_id) as path:
pass
# Allow for backend-specific inspection.
if post_delete_cb is not None:
post_delete_cb(storage_id)
# Again, using upload/download instead of store_path/restore_path.
checkpoints = []
for _ in range(2):
storage_id = str(uuid.uuid4())
path = pathlib.Path(f"/tmp/storage_lifecycle_test-{storage_id}")
try:
create_checkpoint(path)
manager.upload(path, storage_id)
checkpoints.append(storage_id)
finally:
shutil.rmtree(path, ignore_errors=True)
for storage_id in checkpoints:
path = pathlib.Path(f"/tmp/storage_lifecycle_test-{storage_id}")
try:
manager.download(storage_id, path)
validate_checkpoint(path)
finally:
shutil.rmtree(path, ignore_errors=True)
manager.delete(storage_id)
with pytest.raises(errors.CheckpointNotFound):
manager.download(storage_id, path)
if post_delete_cb is not None:
post_delete_cb(storage_id)
def maybe_load_checkpoint(
storage_mgr: storage.StorageManager,
checkpoint: Optional[Dict[str,
Any]]) -> Iterator[Optional[pathlib.Path]]:
"""
Either wrap a storage_mgr.restore_path() context manager, or be a noop
context manager if there is no checkpoint to load.
"""
if checkpoint is None:
yield None
else:
metadata = storage.StorageMetadata.from_json(checkpoint)
logging.info("Restoring trial from checkpoint {}".format(
metadata.storage_id))
with storage_mgr.restore_path(metadata) as path:
yield pathlib.Path(path)