Load only requested demos during fallback conversion

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
 

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()

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)