fix(asset): version marker to redownload assets when new version detected

README.md

@@ -161,7 +161,7 @@ pip install -ve . --no-build-isolation
 
 RCS resolves its asset directory from the `RCS_PREFIX` environment variable. When it is unset, RCS defaults to `~/.rcs`.
 
-On import, RCS checks whether that path exists. If it does not, it downloads the matching asset archive from GitHub into that location automatically.
+On import, RCS checks whether that path contains the assets of the installed RCS version (recorded in `assets/.version`). If not, it downloads the matching asset archive from GitHub into that location automatically, replacing assets of other versions. An `RCS_PREFIX` with an `assets` directory but without a version file, e.g. a source checkout, is used as it is.
 
 ```shell
 export RCS_PREFIX=/path/to/rcs-assets

docs/getting_started/index.md

@@ -46,7 +46,7 @@ For a docker deployment, see the `docker` folder in the repository.
 
 RCS resolves its asset directory from the `RCS_PREFIX` environment variable. If `RCS_PREFIX` is not set, it defaults to `~/.rcs`.
 
-On import, RCS checks whether that path exists. If it does not, it downloads the matching asset archive from GitHub into that location automatically.
+On import, RCS checks whether that path contains the assets of the installed RCS version (recorded in `assets/.version`). If not, it downloads the matching asset archive from GitHub into that location automatically, replacing assets of other versions. An `RCS_PREFIX` with an `assets` directory but without a version file, e.g. a source checkout, is used as it is.
 
 ```shell
 export RCS_PREFIX=/path/to/rcs-assets

python/rcs/__init__.py

@@ -1,7 +1,10 @@
 """Robot control stack python bindings."""
 
+import contextlib
 import io
 import os
+import shutil
+import tempfile
 import zipfile
 from collections.abc import Callable
 from dataclasses import dataclass
@@ -15,7 +18,8 @@
 from rcs import camera, envs, hand, sim
 
 GITHUB_ASSET_ARCHIVE_URL = "https://github.com/RobotControlStack/robot-control-stack/archive/refs/tags/{tag}.zip"
-REQUIRED_ASSET = Path("assets/scenes/empty_world/scene.xml")
+# written into the assets directory after a complete download, holds the version the assets belong to
+ASSETS_VERSION_FILE = ".version"
 
 
 def download_assets(
@@ -28,24 +32,48 @@ def download_assets(
     response = requests.get(github_asset_archive_url.format(tag=tag), timeout=60)
     response.raise_for_status()
 
-    found_assets = False
-    with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
-        for member in archive.infolist():
-            parts = PurePosixPath(member.filename).parts
-            if len(parts) < 2 or parts[1] != assets_dirname:
-                continue
-
-            found_assets = True
-            target = prefix.joinpath(*parts[1:])
-            if member.is_dir():
-                target.mkdir(parents=True, exist_ok=True)
-            else:
-                target.parent.mkdir(parents=True, exist_ok=True)
-                target.write_bytes(archive.read(member))
-
-    if not found_assets:
-        msg = f"Could not find {assets_dirname} in the downloaded GitHub archive."
-        raise FileNotFoundError(msg)
+    prefix.mkdir(parents=True, exist_ok=True)
+    # extract next to the final location and move it into place at the end, so that other processes
+    # never see a partially written assets directory
+    staging = Path(tempfile.mkdtemp(prefix=f".{assets_dirname}-", dir=prefix))
+    try:
+        found_assets = False
+        with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
+            for member in archive.infolist():
+                parts = PurePosixPath(member.filename).parts
+                if len(parts) < 2 or parts[1] != assets_dirname:
+                    continue
+
+                found_assets = True
+                target = staging.joinpath(*parts[1:])
+                if member.is_dir():
+                    target.mkdir(parents=True, exist_ok=True)
+                else:
+                    target.parent.mkdir(parents=True, exist_ok=True)
+                    target.write_bytes(archive.read(member))
+
+        if not found_assets:
+            msg = f"Could not find {assets_dirname} in the downloaded GitHub archive."
+            raise FileNotFoundError(msg)
+
+        new_assets_dir = staging / assets_dirname
+        (new_assets_dir / ASSETS_VERSION_FILE).write_text(version)
+        assets_dir = prefix / assets_dirname
+        with contextlib.suppress(FileNotFoundError):
+            os.replace(assets_dir, staging / "outdated")
+        try:
+            os.replace(new_assets_dir, assets_dir)
+        except OSError:
+            # another process installed the assets in the meantime
+            if not _has_assets_version(assets_dir, version):
+                raise
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
+
+
+def _has_assets_version(assets_dir: Path, version: str) -> bool:
+    version_file = assets_dir / ASSETS_VERSION_FILE
+    return version_file.is_file() and version_file.read_text() == version
 
 
 def get_prefix(
@@ -58,7 +86,14 @@ def get_prefix(
     download_fn: Callable[[str, Path, str, str], None] = download_assets,
 ) -> str:
     def has_required_assets(prefix: Path) -> bool:
-        return (prefix / REQUIRED_ASSET).is_file()
+        assets_dir = prefix / assets_dirname
+        if not assets_dir.is_dir():
+            return False
+        if (assets_dir / ASSETS_VERSION_FILE).is_file():
+            return _has_assets_version(assets_dir, version)
+        # assets without a version file at a user provided prefix (e.g. a source checkout) are used as they are,
+        # in the default prefix they are from an older or interrupted download
+        return bool(env_prefix)
 
     if env_prefix:
         prefix = Path(env_prefix).expanduser().resolve()
@@ -68,10 +103,9 @@ def has_required_assets(prefix: Path) -> bool:
             return str(editable_prefix)
         prefix = default_prefix.resolve()
 
-    assets_dir = prefix / assets_dirname
-    if not assets_dir.is_dir() or not has_required_assets(prefix):
+    if not has_required_assets(prefix):
         prefix.mkdir(parents=True, exist_ok=True)
-        print(f"Assets not found at {assets_dir}, downloading them now.")
+        print(f"Assets for version {version} not found at {prefix / assets_dirname}, downloading them now.")
         download_fn(version, prefix, github_asset_archive_url, assets_dirname)
 
     return str(prefix)

python/tests/test_assets.py

@@ -0,0 +1,105 @@
+import io
+import zipfile
+from pathlib import Path
+
+import pytest
+
+import rcs
+
+URL = "https://example.com/{tag}.zip"
+
+
+class _FakeResponse:
+    def __init__(self, content: bytes):
+        self.content = content
+
+    def raise_for_status(self):
+        pass
+
+
+def _archive(files: dict[str, str]) -> bytes:
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, "w") as archive:
+        for name, content in files.items():
+            archive.writestr(f"robot-control-stack-1.0.0/{name}", content)
+    return buffer.getvalue()
+
+
+@pytest.fixture()
+def downloads(monkeypatch: pytest.MonkeyPatch) -> list[str]:
+    """Serves an archive with a single scene file and records the requested urls."""
+    urls: list[str] = []
+
+    def fake_get(url: str, timeout: float) -> _FakeResponse:  # noqa: ARG001
+        urls.append(url)
+        return _FakeResponse(_archive({"assets/scene.xml": "<mujoco/>", "README.md": ""}))
+
+    monkeypatch.setattr(rcs.requests, "get", fake_get)
+    return urls
+
+
+def _get_prefix(env_prefix: Path | None, tmp_path: Path, version: str = "1.0.0") -> str:
+    return rcs.get_prefix(str(env_prefix) if env_prefix else None, tmp_path / "pkg", tmp_path / "default", version, URL)
+
+
+def test_assets_downloaded_once(tmp_path: Path, downloads: list[str]):
+    for _ in range(3):
+        assert _get_prefix(None, tmp_path) == str(tmp_path / "default")
+    assert downloads == ["https://example.com/v1.0.0.zip"]
+    assets_dir = tmp_path / "default" / "assets"
+    assert (assets_dir / "scene.xml").read_text() == "<mujoco/>"
+    assert (assets_dir / rcs.ASSETS_VERSION_FILE).read_text() == "1.0.0"
+    assert not (tmp_path / "default" / "README.md").exists()
+    # the staging directory is cleaned up
+    assert [p.name for p in (tmp_path / "default").iterdir()] == ["assets"]
+
+
+def test_assets_replaced_on_version_change(tmp_path: Path, downloads: list[str]):
+    _get_prefix(None, tmp_path, version="1.0.0")
+    stale_file = tmp_path / "default" / "assets" / "removed_in_new_version.xml"
+    stale_file.write_text("")
+    _get_prefix(None, tmp_path, version="1.1.0")
+    assert downloads == ["https://example.com/v1.0.0.zip", "https://example.com/v1.1.0.zip"]
+    assert (tmp_path / "default" / "assets" / rcs.ASSETS_VERSION_FILE).read_text() == "1.1.0"
+    assert not stale_file.exists()
+
+
+def test_unversioned_default_assets_redownloaded(tmp_path: Path, downloads: list[str]):
+    # e.g. left by an older rcs version or an interrupted download
+    (tmp_path / "default" / "assets").mkdir(parents=True)
+    _get_prefix(None, tmp_path)
+    assert len(downloads) == 1
+    assert (tmp_path / "default" / "assets" / rcs.ASSETS_VERSION_FILE).is_file()
+
+
+def test_unversioned_env_prefix_used_as_is(tmp_path: Path, downloads: list[str]):
+    # e.g. RCS_PREFIX pointing to a source checkout
+    checkout = tmp_path / "checkout"
+    (checkout / "assets").mkdir(parents=True)
+    assert _get_prefix(checkout, tmp_path) == str(checkout)
+    assert downloads == []
+    assert not (checkout / "assets" / rcs.ASSETS_VERSION_FILE).exists()
+
+
+def test_env_prefix_downloaded_when_missing(tmp_path: Path, downloads: list[str]):
+    prefix = tmp_path / "env"
+    for _ in range(2):
+        assert _get_prefix(prefix, tmp_path) == str(prefix)
+    assert len(downloads) == 1
+
+
+def test_editable_install_skips_download(tmp_path: Path, downloads: list[str]):
+    (tmp_path / "pkg" / "assets").mkdir(parents=True)
+    assert _get_prefix(None, tmp_path) == str(tmp_path / "pkg")
+    assert downloads == []
+
+
+def test_archive_without_assets_keeps_existing_assets(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    assets_dir = tmp_path / "prefix" / "assets"
+    assets_dir.mkdir(parents=True)
+    (assets_dir / "scene.xml").write_text("old")
+    monkeypatch.setattr(rcs.requests, "get", lambda *_, **__: _FakeResponse(_archive({"README.md": ""})))
+    with pytest.raises(FileNotFoundError):
+        rcs.download_assets("1.0.0", tmp_path / "prefix", URL)
+    assert (assets_dir / "scene.xml").read_text() == "old"
+    assert [p.name for p in (tmp_path / "prefix").iterdir()] == ["assets"]