Download model files again after cache eviction
This commit is contained in:
parent
f629b16cd2
commit
63b41faad6
2 changed files with 31 additions and 4 deletions
|
|
@ -4,7 +4,6 @@ import shutil
|
||||||
import sys
|
import sys
|
||||||
import tempfile
|
import tempfile
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from functools import lru_cache
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from types import TracebackType
|
from types import TracebackType
|
||||||
from typing import IO, Any, List, Literal, Optional, Type, Union, cast
|
from typing import IO, Any, List, Literal, Optional, Type, Union, cast
|
||||||
|
|
@ -163,7 +162,6 @@ class LargeFileBase(ABC):
|
||||||
|
|
||||||
return cast(int, self._version)
|
return cast(int, self._version)
|
||||||
|
|
||||||
@lru_cache(1)
|
|
||||||
def get(self, hide_progress: bool = False) -> Path:
|
def get(self, hide_progress: bool = False) -> Path:
|
||||||
"""Return path to the proxy of a file (or directory).
|
"""Return path to the proxy of a file (or directory).
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,8 +1,11 @@
|
||||||
|
import io
|
||||||
|
import tarfile
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Iterator
|
from typing import Any, Iterator
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from great_ai.large_file import LargeFileBase, LargeFileLocal
|
from great_ai.large_file import LargeFileBase, LargeFileLocal, LargeFileS3
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(autouse=True)
|
@pytest.fixture(autouse=True)
|
||||||
|
|
@ -71,3 +74,29 @@ def test_failed_write_removes_temporary_file() -> None:
|
||||||
finally:
|
finally:
|
||||||
if temporary is not None:
|
if temporary is not None:
|
||||||
temporary.unlink(missing_ok=True)
|
temporary.unlink(missing_ok=True)
|
||||||
|
|
||||||
|
|
||||||
|
def test_evicted_remote_model_is_downloaded_again() -> None:
|
||||||
|
backend = MagicMock()
|
||||||
|
backend.list_objects_v2.return_value = {"Contents": [{"Key": "model/0"}]}
|
||||||
|
downloads = []
|
||||||
|
|
||||||
|
def download(
|
||||||
|
self: LargeFileS3, remote_path: Any, local_path: Path, hide_progress: bool
|
||||||
|
) -> None:
|
||||||
|
downloads.append(remote_path)
|
||||||
|
with tarfile.open(local_path, "w:gz") as archive:
|
||||||
|
payload = b"remote weights"
|
||||||
|
member = tarfile.TarInfo("model-0")
|
||||||
|
member.size = len(payload)
|
||||||
|
archive.addfile(member, io.BytesIO(payload))
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(LargeFileS3, "_client", backend),
|
||||||
|
patch.object(LargeFileS3, "_download", download),
|
||||||
|
):
|
||||||
|
reader = LargeFileS3("model")
|
||||||
|
cached = reader.get()
|
||||||
|
cached.unlink()
|
||||||
|
assert reader.get().read_text() == "remote weights"
|
||||||
|
assert downloads == ["model/0", "model/0"]
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue