diff --git a/dimos/utils/data.py b/dimos/utils/data.py index b4b7e2210a..d5b7f7aa59 100644 --- a/dimos/utils/data.py +++ b/dimos/utils/data.py @@ -14,17 +14,21 @@ from datetime import datetime from functools import cache +import json import os from pathlib import Path import platform import re +import shutil import subprocess import sys import tarfile import tempfile import time -from dimos.constants import DIMOS_PROJECT_ROOT +from filelock import FileLock + +from dimos.constants import DIMOS_PROJECT_ROOT, STATE_DIR from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -221,9 +225,89 @@ def _lfs_pull(file_path: Path, repo_root: Path, *, retries: int = 2) -> None: def _decompress_archive(filename: str | Path) -> Path: target_dir = get_data_dir() filename_path = Path(filename) - with tarfile.open(filename_path, "r:gz") as tar: - tar.extractall(target_dir) - return target_dir / filename_path.name.replace(".tar.gz", "") + extracted_name = filename_path.name.replace(".tar.gz", "") + extracted_path = target_dir / extracted_name + + # extract the file to a temp location first, then swap if that worked + staging_dir = Path(tempfile.mkdtemp(dir=target_dir, prefix=f".{extracted_name}.staging-")) + try: + with tarfile.open(filename_path, "r:gz") as tar: + tar.extractall(staging_dir) + + # validate before touching the old extraction, so a mis-packaged + # archive can't delete a good copy and then fail to replace it + staged_path = staging_dir / extracted_name + if not staged_path.exists(): + raise RuntimeError( + f"Archive '{filename_path.name}' has no top-level '{extracted_name}' member" + ) + + if extracted_path.is_dir(): + shutil.rmtree(extracted_path) + elif extracted_path.exists(): + extracted_path.unlink() + os.replace(staged_path, extracted_path) + finally: + shutil.rmtree(staging_dir, ignore_errors=True) + + return extracted_path + + +def _lfs_archive_path(archive_name: str | Path) -> Path: + return _get_lfs_dir() / (str(archive_name) + ".tar.gz") + + +@cache +def _lfs_extraction_state_dir() -> Path: + """Per-checkout directory for extraction stamps and locks.""" + # Keyed by the data dir's own path so worktrees don't collide. + checkout_key = str(get_data_dir().resolve()).replace("/", "_").lstrip("_") + return STATE_DIR / "lfs-extractions" / checkout_key + + +def _extraction_stamp_path(archive_name: str | Path) -> Path: + return _lfs_extraction_state_dir() / f"{archive_name}.json" + + +def _archive_lock_path(archive_name: str | Path) -> Path: + return _lfs_extraction_state_dir() / f"{archive_name}.lock" + + +def _archive_lock(archive_name: str | Path) -> FileLock: + _lfs_extraction_state_dir().mkdir(parents=True, exist_ok=True) + return FileLock(_archive_lock_path(archive_name)) + + +def _archive_identity(archive_path: Path) -> dict[str, int]: + stat_result = archive_path.stat() + return {"st_size": stat_result.st_size, "st_mtime_ns": stat_result.st_mtime_ns} + + +def _extraction_is_current(archive_name: str | Path) -> bool: + """Whether the extracted data still matches its archive. + + No archive behind it means a local recording: always current. + """ + archive_path = _lfs_archive_path(archive_name) + if not archive_path.exists(): + return True + + stamp_path = _extraction_stamp_path(archive_name) + if not stamp_path.exists(): + return False + + try: + stamp: dict[str, int] = json.loads(stamp_path.read_text()) + return stamp == _archive_identity(archive_path) + except (json.JSONDecodeError, OSError): + # if we can't load the stamp, just treat it as stale + return False + + +def _write_extraction_stamp(archive_name: str | Path, archive_path: Path) -> None: + stamp_path = _extraction_stamp_path(archive_name) + stamp_path.parent.mkdir(parents=True, exist_ok=True) + stamp_path.write_text(json.dumps(_archive_identity(archive_path))) def _pull_lfs_archive(filename: str | Path) -> Path: @@ -234,7 +318,7 @@ def _pull_lfs_archive(filename: str | Path) -> Path: repo_root = get_project_root() # Construct path to test data file - file_path = _get_lfs_dir() / (str(filename) + ".tar.gz") + file_path = _lfs_archive_path(filename) # Check if file exists if not file_path.exists(): @@ -291,17 +375,25 @@ def get_data(name: str | Path) -> Path: data_dir = get_data_dir() file_path = data_dir / name - # already pulled and decompressed, return it directly - if file_path.exists(): - return file_path - # extract archive root (first path component) and nested path path_parts = Path(name).parts archive_name = path_parts[0] nested_path = Path(*path_parts[1:]) if len(path_parts) > 1 else None - # download and decompress the archive root - archive_path = _decompress_archive(_pull_lfs_archive(archive_name)) + def is_ready() -> bool: + # already extracted and current + return file_path.exists() and _extraction_is_current(archive_name) + + if is_ready(): + return file_path + + with _archive_lock(archive_name): + if not is_ready(): + archive_tar_path = _pull_lfs_archive(archive_name) + archive_path = _decompress_archive(archive_tar_path) + _write_extraction_stamp(archive_name, archive_tar_path) + else: + archive_path = data_dir / archive_name # return full path including nested components if nested_path: diff --git a/dimos/utils/test_data.py b/dimos/utils/test_data.py index 7564d4dfd0..8783a3d1a6 100644 --- a/dimos/utils/test_data.py +++ b/dimos/utils/test_data.py @@ -12,10 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass import hashlib +import io +import json import os from pathlib import Path import subprocess +import tarfile import pytest @@ -100,6 +106,222 @@ def test_backup_file_keep_last_zero_removes_all(tmp_path: Path) -> None: assert list(tmp_path.glob("recording_go2.*.db")) == [] +@dataclass +class DataLayout: + data_dir: Path + lfs_dir: Path + state_dir: Path + staging_root: Path + + +@pytest.fixture +def data_layout(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> DataLayout: + layout = DataLayout( + data_dir=tmp_path / "data", + lfs_dir=tmp_path / "data" / ".lfs", + state_dir=tmp_path / "state", + staging_root=tmp_path / "staging", + ) + layout.data_dir.mkdir() + layout.lfs_dir.mkdir() + layout.staging_root.mkdir() + + monkeypatch.setattr( + data, + "get_data_dir", + lambda extra_path=None: layout.data_dir / extra_path if extra_path else layout.data_dir, + ) + monkeypatch.setattr(data, "_get_lfs_dir", lambda: layout.lfs_dir) + monkeypatch.setattr(data, "_lfs_extraction_state_dir", lambda: layout.state_dir) + monkeypatch.setattr(data, "_check_git_lfs_available", lambda: True) + return layout + + +def _write_tar_gz(tar_path: Path, archive_name: str, content: bytes, staging_root: Path) -> None: + staging = staging_root / archive_name + staging.mkdir(exist_ok=True) + (staging / "payload.bin").write_bytes(content) + with tarfile.open(tar_path, "w:gz") as tar: + tar.add(staging, arcname=archive_name) + + +def _write_single_file_tar_gz(tar_path: Path, member_name: str, content: bytes) -> None: + with tarfile.open(tar_path, "w:gz") as tar: + info = tarfile.TarInfo(member_name) + info.size = len(content) + tar.addfile(info, io.BytesIO(content)) + + +def _counting_decompress( + monkeypatch: pytest.MonkeyPatch, +) -> Callable[[], int]: + """Wrap _decompress_archive and return its running call count.""" + call_count = 0 + original_decompress = data._decompress_archive + + def counting_decompress(filename: str | Path) -> Path: + nonlocal call_count + call_count += 1 + return original_decompress(filename) + + monkeypatch.setattr(data, "_decompress_archive", counting_decompress) + return lambda: call_count + + +def test_get_data_self_heals_legacy_extraction_without_stamp(data_layout: DataLayout) -> None: + """No stamp file means not trusted: re-extracts once, then self-heals.""" + tar_path = data._lfs_archive_path("dataset") + _write_tar_gz(tar_path, "dataset", b"current", data_layout.staging_root) + + # archive and extraction both present, but no stamp + (data_layout.data_dir / "dataset").mkdir() + (data_layout.data_dir / "dataset" / "payload.bin").write_bytes(b"stale") + + healed = data.get_data("dataset/payload.bin") + assert healed.read_bytes() == b"current" + assert data._extraction_stamp_path("dataset").exists() + + +def test_get_data_returns_local_recording_without_backing_archive( + data_layout: DataLayout, monkeypatch: pytest.MonkeyPatch +) -> None: + local_db = data_layout.data_dir / "recording.db" + local_db.write_bytes(b"local only, no archive behind it") + + def fail_pull(name: str) -> Path: + raise AssertionError("must not attempt an LFS pull for a local-only recording") + + monkeypatch.setattr(data, "_pull_lfs_archive", fail_pull) + + assert data.get_data("recording.db") == local_db + + +def test_get_data_reextracts_when_archive_changes_then_stabilizes( + data_layout: DataLayout, monkeypatch: pytest.MonkeyPatch +) -> None: + tar_path = data._lfs_archive_path("dataset") + _write_tar_gz(tar_path, "dataset", b"v1", data_layout.staging_root) + + first = data.get_data("dataset/payload.bin") + assert first.read_bytes() == b"v1" + + call_count = _counting_decompress(monkeypatch) + + # unchanged archive: no re-extraction + data.get_data("dataset/payload.bin") + assert call_count() == 0 + + # archive re-tarred under the same name with different content + _write_tar_gz( + tar_path, "dataset", b"v2, a longer payload than before", data_layout.staging_root + ) + + second = data.get_data("dataset/payload.bin") + assert second.read_bytes() == b"v2, a longer payload than before" + assert call_count() == 1 + + # settles back to no re-extraction once the stamp matches again + data.get_data("dataset/payload.bin") + assert call_count() == 1 + + +def test_get_data_reextraction_drops_members_removed_from_archive(data_layout: DataLayout) -> None: + """Members dropped from a re-tarred archive must not survive on disk.""" + staging = data_layout.staging_root / "dataset" + staging.mkdir(parents=True) + tar_path = data._lfs_archive_path("dataset") + + (staging / "keep.bin").write_bytes(b"kept") + (staging / "drop.bin").write_bytes(b"dropped") + with tarfile.open(tar_path, "w:gz") as tar: + tar.add(staging, arcname="dataset") + + dropped_path = data.get_data("dataset/drop.bin") + assert dropped_path.exists() + + # re-tar without drop.bin + (staging / "drop.bin").unlink() + with tarfile.open(tar_path, "w:gz") as tar: + tar.add(staging, arcname="dataset") + + kept_path = data.get_data("dataset/keep.bin") + assert kept_path.read_bytes() == b"kept" + assert not dropped_path.exists() + + +def test_get_data_mispackaged_archive_does_not_destroy_prior_extraction( + data_layout: DataLayout, +) -> None: + """A stale-triggering archive missing its top-level member must raise + without deleting the still-good previous extraction.""" + tar_path = data._lfs_archive_path("dataset") + _write_tar_gz(tar_path, "dataset", b"v1", data_layout.staging_root) + + first = data.get_data("dataset/payload.bin") + assert first.read_bytes() == b"v1" + + # re-tar under the wrong top-level name, so extraction succeeds but the + # expected "dataset" member never lands in the staging dir + with tarfile.open(tar_path, "w:gz") as tar: + info = tarfile.TarInfo("wrong_name/payload.bin") + payload = b"v2" + info.size = len(payload) + tar.addfile(info, io.BytesIO(payload)) + + with pytest.raises(RuntimeError, match="dataset"): + data.get_data("dataset/payload.bin") + + assert first.exists() + assert first.read_bytes() == b"v1" + + +def test_get_data_reextracts_when_stamp_is_corrupt(data_layout: DataLayout) -> None: + tar_path = data._lfs_archive_path("dataset") + _write_tar_gz(tar_path, "dataset", b"current", data_layout.staging_root) + + stamp_path = data._extraction_stamp_path("dataset") + stamp_path.parent.mkdir(parents=True, exist_ok=True) + stamp_path.write_text("not valid json") + (data_layout.data_dir / "dataset").mkdir() + (data_layout.data_dir / "dataset" / "payload.bin").write_bytes(b"stale") + + result = data.get_data("dataset/payload.bin") + + assert result.read_bytes() == b"current" + assert json.loads(stamp_path.read_text()) + + +def test_get_data_reextracts_single_file_archive(data_layout: DataLayout) -> None: + """A prior extraction that is a bare file, not a directory, must also be replaced.""" + tar_path = data._lfs_archive_path("cafe.jpg") + _write_single_file_tar_gz(tar_path, "cafe.jpg", b"v1") + + first = data.get_data("cafe.jpg") + assert first.is_file() + assert first.read_bytes() == b"v1" + + _write_single_file_tar_gz(tar_path, "cafe.jpg", b"v2, a longer payload than before") + + second = data.get_data("cafe.jpg") + assert second.read_bytes() == b"v2, a longer payload than before" + + +def test_get_data_serializes_concurrent_extraction_of_same_archive( + data_layout: DataLayout, monkeypatch: pytest.MonkeyPatch +) -> None: + """Concurrent callers must not race: only one extraction happens.""" + tar_path = data._lfs_archive_path("dataset") + _write_tar_gz(tar_path, "dataset", b"v1", data_layout.staging_root) + + call_count = _counting_decompress(monkeypatch) + + with ThreadPoolExecutor(max_workers=8) as executor: + results = list(executor.map(lambda _: data.get_data("dataset/payload.bin"), range(8))) + + assert all(result.read_bytes() == b"v1" for result in results) + assert call_count() == 1 + + @pytest.mark.self_hosted def test_pull_file() -> None: repo_root = data.get_project_root()