Download model files again after cache eviction

This commit is contained in:
Andras Schmelczer 2026-10-27 16:53:50 +00:00
commit 63b41faad6
2 changed files with 31 additions and 4 deletions

View file

@ -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).

View file

@ -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"]