From f7309da8f9bbd18e4141776e0c8bceb316a9c033 Mon Sep 17 00:00:00 2001 From: eust-w Date: Thu, 6 Aug 2026 05:41:41 +0800 Subject: [PATCH] Load only requested demos during fallback --- CHANGELOG.md | 3 ++- demonstrations/demo_store.py | 27 ++++++++++++++++++--------- tests/test_demo_store.py | 22 ++++++++++++++++++++++ 3 files changed, 42 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2230e87..092aeeb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed -- None. +- Demo fallback conversion now loads only the requested number of safetensors + files instead of materialising the complete task archive first. ## 4.2.0 diff --git a/demonstrations/demo_store.py b/demonstrations/demo_store.py index 4b2d9a1..ab597f6 100644 --- a/demonstrations/demo_store.py +++ b/demonstrations/demo_store.py @@ -121,28 +121,30 @@ def get_demos( if demos: return demos - # Attempt to get demos at original frequency + # Attempt to get demos at original frequency. Keep paths here rather + # than materialising every safetensors file before applying ``amount``. max_freq_demos_dir = self._create_path(metadata).parent - demos = self._get_demos(max_freq_demos_dir, -1) + demo_paths = self._get_demo_paths(max_freq_demos_dir, amount) # Attempt to get raw, lightweight demos - if not demos and metadata.observation_mode != ObservationMode.Lightweight: - demos = self._get_demos(light_demos_dir, -1) + if not demo_paths and metadata.observation_mode != ObservationMode.Lightweight: + demo_paths = self._get_demo_paths(light_demos_dir, amount) # Raising exception if there are no demos at this stage - if not demos: + if not demo_paths: raise DemoNotFoundError(metadata) # Recreate and cache demos with tqdm( - total=len(demos), + total=len(demo_paths), desc="Caching Demos", unit="demo", leave=True, ) as pbar: robot = metadata.get_robot() env = metadata.get_env(frequency) - for demo in demos: + for demo_path in demo_paths: + demo = Demo.from_safetensors(demo_path) if self.demo_exists(metadata, frequency, demo.uuid): pbar.update() continue @@ -158,7 +160,8 @@ def get_demos( pbar.update() return self._get_demos(demos_dir, amount) - def _get_demos(self, demos_dir: Path, amount: int) -> list[Demo]: + def _get_demo_paths(self, demos_dir: Path, amount: int) -> list[Path]: + """Return shuffled demo paths, limited before deserialisation.""" self.pull_demos() if not demos_dir.exists(): return [] @@ -169,7 +172,13 @@ def _get_demos(self, demos_dir: Path, amount: int) -> list[Demo]: np.random.shuffle(files) if amount > 0: files = files[:amount] - return [Demo.from_safetensors(file) for file in files] + return files + + def _get_demos(self, demos_dir: Path, amount: int) -> list[Demo]: + return [ + Demo.from_safetensors(file) + for file in self._get_demo_paths(demos_dir, amount) + ] def _get_demos_count(self, demos_dir: Path) -> int: self.pull_demos() diff --git a/tests/test_demo_store.py b/tests/test_demo_store.py index 39c609d..3bed320 100644 --- a/tests/test_demo_store.py +++ b/tests/test_demo_store.py @@ -1,6 +1,7 @@ """Test demonstration storing system.""" import tempfile +import numpy as np import pytest from pathlib import Path from typing import Optional @@ -311,6 +312,27 @@ def test_retrieve_n_demos(self, temp_demo_store): demos_from_store = temp_demo_store.get_demos(metadata, amount=i) assert len(demos_from_store) == i + def test_loads_only_requested_demo_files( + self, temp_demo_store, monkeypatch, tmp_path + ): + demos_dir = tmp_path / "source-demos" + demos_dir.mkdir() + for index in range(5): + (demos_dir / f"{index}{SAFETENSORS_SUFFIX}").touch() + + loaded_paths = [] + monkeypatch.setattr(np.random, "shuffle", lambda _paths: None) + monkeypatch.setattr( + Demo, + "from_safetensors", + lambda path: loaded_paths.append(path) or path, + ) + + demos = temp_demo_store._get_demos(demos_dir, amount=2) + + assert demos == loaded_paths + assert len(loaded_paths) == 2 + def test_implicit_saving_of_lightweight_demos(self, temp_demo_store): demo = _generate_simple_demo() temp_demo_store.cache_demo(demo)