mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 12:15:59 +00:00
Harden production dataset streaming pipeline
This commit is contained in:
@@ -178,7 +178,9 @@ To capture comparable data-pipeline results on a training host:
|
|||||||
uv run python scripts/bench_streaming_dataset.py \
|
uv run python scripts/bench_streaming_dataset.py \
|
||||||
--repo-id=yaak-ai/L2D-v3 \
|
--repo-id=yaak-ai/L2D-v3 \
|
||||||
--batch-size=16 \
|
--batch-size=16 \
|
||||||
--num-workers=4 \
|
--num-workers=1 \
|
||||||
|
--fetch-workers=4 \
|
||||||
|
--decode-threads=2 \
|
||||||
--summary-json=streaming-benchmark.json
|
--summary-json=streaming-benchmark.json
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -58,12 +58,17 @@ Subset sidecars cannot be published.
|
|||||||
|
|
||||||
The production defaults are:
|
The production defaults are:
|
||||||
|
|
||||||
| Option | Default | Meaning |
|
| Option | Default | Meaning |
|
||||||
| ----------------------------- | ------: | --------------------------------------------------- |
|
| ----------------------------------- | ------: | --------------------------------------------------- |
|
||||||
| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank |
|
| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank |
|
||||||
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
|
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
|
||||||
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank |
|
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank |
|
||||||
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers |
|
||||||
|
| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order |
|
||||||
|
| `streaming_max_open_decoders` | 64 | Independent open-decoder LRU cap per rank |
|
||||||
|
| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank |
|
||||||
|
| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read |
|
||||||
|
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
||||||
|
|
||||||
The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
|
The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes
|
||||||
computed from the sidecar. An episode larger than the complete rank budget fails before training
|
computed from the sidecar. An episode larger than the complete rank budget fails before training
|
||||||
@@ -77,6 +82,9 @@ lerobot-train \
|
|||||||
--dataset.streaming_episode_pool_size=16 \
|
--dataset.streaming_episode_pool_size=16 \
|
||||||
--dataset.streaming_prefetch_episodes=4 \
|
--dataset.streaming_prefetch_episodes=4 \
|
||||||
--dataset.streaming_byte_budget_gb=4 \
|
--dataset.streaming_byte_budget_gb=4 \
|
||||||
|
--dataset.streaming_decode_threads=2 \
|
||||||
|
--dataset.streaming_decoded_queue_size=8 \
|
||||||
|
--dataset.streaming_max_open_decoders=64 \
|
||||||
--num_workers=4 \
|
--num_workers=4 \
|
||||||
--policy.type=act \
|
--policy.type=act \
|
||||||
--output_dir=outputs/train/act_streaming
|
--output_dir=outputs/train/act_streaming
|
||||||
@@ -109,6 +117,8 @@ uv run python scripts/bench_streaming_dataset.py \
|
|||||||
--batch-size=16 \
|
--batch-size=16 \
|
||||||
--num-workers=1 \
|
--num-workers=1 \
|
||||||
--fetch-workers=4 \
|
--fetch-workers=4 \
|
||||||
|
--decode-threads=2 \
|
||||||
|
--decoded-queue-size=8 \
|
||||||
--summary-json=streaming-benchmark.json
|
--summary-json=streaming-benchmark.json
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -70,7 +70,7 @@ def parse_args() -> argparse.Namespace:
|
|||||||
type=int,
|
type=int,
|
||||||
default=8,
|
default=8,
|
||||||
help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; "
|
help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; "
|
||||||
"the HF bucket path saturates around 64 connections per host, so keep the product near 64.",
|
"measure the host-specific saturation point before increasing this value.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--range-subranges",
|
"--range-subranges",
|
||||||
@@ -443,6 +443,7 @@ def run_exact_coverage_stream(
|
|||||||
decode_pool.shutdown(wait=True)
|
decode_pool.shutdown(wait=True)
|
||||||
|
|
||||||
elapsed = time.perf_counter() - start
|
elapsed = time.perf_counter() - start
|
||||||
|
replacements = max(0, pool.admitted_count - min(pool_size, len(order)))
|
||||||
result = {
|
result = {
|
||||||
"target_samples_s": target_samples_s,
|
"target_samples_s": target_samples_s,
|
||||||
"actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"),
|
"actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"),
|
||||||
@@ -453,6 +454,10 @@ def run_exact_coverage_stream(
|
|||||||
"shard_total_frames": float(total_frames),
|
"shard_total_frames": float(total_frames),
|
||||||
"epoch_complete": 1.0 if epoch_complete else 0.0,
|
"epoch_complete": 1.0 if epoch_complete else 0.0,
|
||||||
"prefetch_ahead": float(prefetch_ahead),
|
"prefetch_ahead": float(prefetch_ahead),
|
||||||
|
"prefetch_episodes": float(prefetch_ahead),
|
||||||
|
"replacements": float(replacements),
|
||||||
|
"replacement_episodes_s": replacements / elapsed if elapsed > 0 else float("inf"),
|
||||||
|
"samples_per_episode": samples_done / replacements if replacements else float(samples_done),
|
||||||
"batch_size": float(batch_size),
|
"batch_size": float(batch_size),
|
||||||
"decode_workers": float(decode_workers),
|
"decode_workers": float(decode_workers),
|
||||||
"kept_up": 1.0
|
"kept_up": 1.0
|
||||||
|
|||||||
@@ -50,6 +50,11 @@ def parse_args() -> argparse.Namespace:
|
|||||||
parser.add_argument("--episode-pool-size", type=int, default=32)
|
parser.add_argument("--episode-pool-size", type=int, default=32)
|
||||||
parser.add_argument("--prefetch-episodes", type=int, default=8)
|
parser.add_argument("--prefetch-episodes", type=int, default=8)
|
||||||
parser.add_argument("--byte-budget-gb", type=float, default=8.0)
|
parser.add_argument("--byte-budget-gb", type=float, default=8.0)
|
||||||
|
parser.add_argument("--decode-threads", type=int, default=2)
|
||||||
|
parser.add_argument("--decoded-queue-size", type=int, default=8)
|
||||||
|
parser.add_argument("--max-open-decoders", type=int, default=64)
|
||||||
|
parser.add_argument("--native-http-connections", type=int, default=None)
|
||||||
|
parser.add_argument("--native-http-subranges", type=int, default=1)
|
||||||
parser.add_argument("--warmup-batches", type=int, default=8)
|
parser.add_argument("--warmup-batches", type=int, default=8)
|
||||||
parser.add_argument("--measure-batches", type=int, default=128)
|
parser.add_argument("--measure-batches", type=int, default=128)
|
||||||
parser.add_argument("--summary-json", type=Path, default=None)
|
parser.add_argument("--summary-json", type=Path, default=None)
|
||||||
@@ -105,6 +110,11 @@ def main() -> None:
|
|||||||
episode_pool_size=args.episode_pool_size,
|
episode_pool_size=args.episode_pool_size,
|
||||||
prefetch_episodes=args.prefetch_episodes,
|
prefetch_episodes=args.prefetch_episodes,
|
||||||
byte_budget_gb=args.byte_budget_gb,
|
byte_budget_gb=args.byte_budget_gb,
|
||||||
|
decode_threads=args.decode_threads,
|
||||||
|
decoded_queue_size=args.decoded_queue_size,
|
||||||
|
max_open_decoders=args.max_open_decoders,
|
||||||
|
native_http_connections=args.native_http_connections,
|
||||||
|
native_http_subranges=args.native_http_subranges,
|
||||||
max_num_shards=max(1, args.fetch_workers),
|
max_num_shards=max(1, args.fetch_workers),
|
||||||
return_uint8=True,
|
return_uint8=True,
|
||||||
)
|
)
|
||||||
@@ -184,6 +194,11 @@ def main() -> None:
|
|||||||
"episode_pool_size": args.episode_pool_size,
|
"episode_pool_size": args.episode_pool_size,
|
||||||
"prefetch_episodes": args.prefetch_episodes,
|
"prefetch_episodes": args.prefetch_episodes,
|
||||||
"byte_budget_gb": args.byte_budget_gb,
|
"byte_budget_gb": args.byte_budget_gb,
|
||||||
|
"decode_threads": args.decode_threads,
|
||||||
|
"decoded_queue_size": args.decoded_queue_size,
|
||||||
|
"max_open_decoders": args.max_open_decoders,
|
||||||
|
"native_http_connections": args.native_http_connections,
|
||||||
|
"native_http_subranges": args.native_http_subranges,
|
||||||
"warmup_batches": args.warmup_batches,
|
"warmup_batches": args.warmup_batches,
|
||||||
"measure_batches": args.measure_batches,
|
"measure_batches": args.measure_batches,
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -52,6 +52,14 @@ class DatasetConfig:
|
|||||||
streaming_prefetch_episodes: int = 8
|
streaming_prefetch_episodes: int = 8
|
||||||
# Hard per-rank cap for synthesized episode-video bytes.
|
# Hard per-rank cap for synthesized episode-video bytes.
|
||||||
streaming_byte_budget_gb: float = 8.0
|
streaming_byte_budget_gb: float = 8.0
|
||||||
|
# Parallel sample assembly/decode workers and their bounded in-order result queue.
|
||||||
|
streaming_decode_threads: int = 2
|
||||||
|
streaming_decoded_queue_size: int = 8
|
||||||
|
# Independent decoder-state cap.
|
||||||
|
streaming_max_open_decoders: int = 64
|
||||||
|
# Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
|
||||||
|
streaming_native_http_connections: int | None = None
|
||||||
|
streaming_native_http_subranges: int = 1
|
||||||
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
|
# Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
|
||||||
eval_split: float = 0.0
|
eval_split: float = 0.0
|
||||||
|
|
||||||
@@ -68,6 +76,16 @@ class DatasetConfig:
|
|||||||
raise ValueError("streaming_prefetch_episodes must be non-negative")
|
raise ValueError("streaming_prefetch_episodes must be non-negative")
|
||||||
if self.streaming_byte_budget_gb <= 0:
|
if self.streaming_byte_budget_gb <= 0:
|
||||||
raise ValueError("streaming_byte_budget_gb must be positive")
|
raise ValueError("streaming_byte_budget_gb must be positive")
|
||||||
|
if self.streaming_decode_threads <= 0:
|
||||||
|
raise ValueError("streaming_decode_threads must be positive")
|
||||||
|
if self.streaming_decoded_queue_size <= 0:
|
||||||
|
raise ValueError("streaming_decoded_queue_size must be positive")
|
||||||
|
if self.streaming_max_open_decoders <= 0:
|
||||||
|
raise ValueError("streaming_max_open_decoders must be positive")
|
||||||
|
if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0:
|
||||||
|
raise ValueError("streaming_native_http_connections must be positive")
|
||||||
|
if self.streaming_native_http_subranges <= 0:
|
||||||
|
raise ValueError("streaming_native_http_subranges must be positive")
|
||||||
if self.episodes is not None:
|
if self.episodes is not None:
|
||||||
if any(ep < 0 for ep in self.episodes):
|
if any(ep < 0 for ep in self.episodes):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|||||||
@@ -119,6 +119,11 @@ def make_dataset(
|
|||||||
episode_pool_size=cfg.dataset.streaming_episode_pool_size,
|
episode_pool_size=cfg.dataset.streaming_episode_pool_size,
|
||||||
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
|
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
|
||||||
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb,
|
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb,
|
||||||
|
decode_threads=cfg.dataset.streaming_decode_threads,
|
||||||
|
decoded_queue_size=cfg.dataset.streaming_decoded_queue_size,
|
||||||
|
max_open_decoders=cfg.dataset.streaming_max_open_decoders,
|
||||||
|
native_http_connections=cfg.dataset.streaming_native_http_connections,
|
||||||
|
native_http_subranges=cfg.dataset.streaming_native_http_subranges,
|
||||||
repeat=True,
|
repeat=True,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
@@ -210,6 +215,11 @@ def make_train_eval_datasets(
|
|||||||
episode_pool_size=cfg.dataset.streaming_episode_pool_size,
|
episode_pool_size=cfg.dataset.streaming_episode_pool_size,
|
||||||
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
|
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
|
||||||
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb,
|
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb,
|
||||||
|
decode_threads=cfg.dataset.streaming_decode_threads,
|
||||||
|
decoded_queue_size=cfg.dataset.streaming_decoded_queue_size,
|
||||||
|
max_open_decoders=cfg.dataset.streaming_max_open_decoders,
|
||||||
|
native_http_connections=cfg.dataset.streaming_native_http_connections,
|
||||||
|
native_http_subranges=cfg.dataset.streaming_native_http_subranges,
|
||||||
repeat=True,
|
repeat=True,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -15,6 +15,7 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
import io
|
import io
|
||||||
import os
|
import os
|
||||||
|
from collections import deque
|
||||||
from collections.abc import Callable, Iterator, Mapping
|
from collections.abc import Callable, Iterator, Mapping
|
||||||
from concurrent.futures import Future, ThreadPoolExecutor
|
from concurrent.futures import Future, ThreadPoolExecutor
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -122,6 +123,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
repeat: bool = False,
|
repeat: bool = False,
|
||||||
*,
|
*,
|
||||||
token: str | bool | None = None,
|
token: str | bool | None = None,
|
||||||
|
decode_threads: int = 2,
|
||||||
|
decoded_queue_size: int = 8,
|
||||||
|
max_open_decoders: int = 64,
|
||||||
|
native_http_connections: int | None = None,
|
||||||
|
native_http_subranges: int = 1,
|
||||||
):
|
):
|
||||||
"""Initialize a StreamingLeRobotDataset.
|
"""Initialize a StreamingLeRobotDataset.
|
||||||
|
|
||||||
@@ -152,6 +158,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
episode_pool_size (int | None, optional): Number of complete episodes in the sampling pool.
|
episode_pool_size (int | None, optional): Number of complete episodes in the sampling pool.
|
||||||
prefetch_episodes (int, optional): Episodes prefetched beyond the active pool.
|
prefetch_episodes (int, optional): Episodes prefetched beyond the active pool.
|
||||||
byte_budget_gb (float, optional): Per-rank upper bound for synthesized episode video bytes.
|
byte_budget_gb (float, optional): Per-rank upper bound for synthesized episode video bytes.
|
||||||
|
decode_threads (int, optional): Parallel sample assembly and video decode workers.
|
||||||
|
decoded_queue_size (int, optional): Maximum number of decoded samples produced ahead
|
||||||
|
of the consumer. Results are yielded in exact planner order.
|
||||||
|
max_open_decoders (int, optional): Maximum number of open video decoders per rank.
|
||||||
|
native_http_connections (int | None, optional): Native HTTP connection limit per rank.
|
||||||
|
``None`` preserves the fetcher's worker-derived default.
|
||||||
|
native_http_subranges (int, optional): Concurrent HTTP subranges per video range read.
|
||||||
repeat (bool, optional): Repeat rank-local exact-coverage epochs without yielding a
|
repeat (bool, optional): Repeat rank-local exact-coverage epochs without yielding a
|
||||||
short final training batch. The training factory enables this; direct iteration is
|
short final training batch. The training factory enables this; direct iteration is
|
||||||
finite by default.
|
finite by default.
|
||||||
@@ -196,9 +209,24 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
raise ValueError("prefetch_episodes must be non-negative")
|
raise ValueError("prefetch_episodes must be non-negative")
|
||||||
if byte_budget_gb <= 0:
|
if byte_budget_gb <= 0:
|
||||||
raise ValueError("byte_budget_gb must be positive")
|
raise ValueError("byte_budget_gb must be positive")
|
||||||
|
if decode_threads <= 0:
|
||||||
|
raise ValueError("decode_threads must be positive")
|
||||||
|
if decoded_queue_size <= 0:
|
||||||
|
raise ValueError("decoded_queue_size must be positive")
|
||||||
|
if max_open_decoders <= 0:
|
||||||
|
raise ValueError("max_open_decoders must be positive")
|
||||||
|
if native_http_connections is not None and native_http_connections <= 0:
|
||||||
|
raise ValueError("native_http_connections must be positive")
|
||||||
|
if native_http_subranges <= 0:
|
||||||
|
raise ValueError("native_http_subranges must be positive")
|
||||||
self.episode_pool_size = episode_pool_size or min(buffer_size, 32)
|
self.episode_pool_size = episode_pool_size or min(buffer_size, 32)
|
||||||
self.prefetch_episodes = prefetch_episodes
|
self.prefetch_episodes = prefetch_episodes
|
||||||
self.byte_budget = int(byte_budget_gb * 1024**3)
|
self.byte_budget = int(byte_budget_gb * 1024**3)
|
||||||
|
self.decode_threads = decode_threads
|
||||||
|
self.decoded_queue_size = decoded_queue_size
|
||||||
|
self.max_open_decoders = max_open_decoders
|
||||||
|
self.native_http_connections = native_http_connections
|
||||||
|
self.native_http_subranges = native_http_subranges
|
||||||
self.repeat = repeat
|
self.repeat = repeat
|
||||||
self._next_epoch = 0
|
self._next_epoch = 0
|
||||||
self._active_epoch = 0
|
self._active_epoch = 0
|
||||||
@@ -356,8 +384,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
columns=self._projected_columns,
|
columns=self._projected_columns,
|
||||||
token=self._streaming_io_token,
|
token=self._streaming_io_token,
|
||||||
)
|
)
|
||||||
executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet")
|
parquet_executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet")
|
||||||
|
decode_executor = ThreadPoolExecutor(
|
||||||
|
max_workers=self.decode_threads,
|
||||||
|
thread_name_prefix="lerobot-decode",
|
||||||
|
)
|
||||||
episode_futures: dict[int, Future[datasets.Dataset]] = {}
|
episode_futures: dict[int, Future[datasets.Dataset]] = {}
|
||||||
|
decoded_futures: deque[Future[dict]] = deque()
|
||||||
scheduled_episodes: set[int] = set()
|
scheduled_episodes: set[int] = set()
|
||||||
retained_video_episodes: set[int] = set()
|
retained_video_episodes: set[int] = set()
|
||||||
if video_cache is not None:
|
if video_cache is not None:
|
||||||
@@ -368,7 +401,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
def submit(episode_index: int) -> Future[datasets.Dataset]:
|
def submit(episode_index: int) -> Future[datasets.Dataset]:
|
||||||
future = episode_futures.get(episode_index)
|
future = episode_futures.get(episode_index)
|
||||||
if future is None:
|
if future is None:
|
||||||
future = executor.submit(self._load_episode_dataset, parquet_reader, episode_index)
|
future = parquet_executor.submit(self._load_episode_dataset, parquet_reader, episode_index)
|
||||||
episode_futures[episode_index] = future
|
episode_futures[episode_index] = future
|
||||||
return future
|
return future
|
||||||
|
|
||||||
@@ -382,44 +415,83 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
video_cache.submit_prefetch(episode_index)
|
video_cache.submit_prefetch(episode_index)
|
||||||
scheduled_episodes.add(episode_index)
|
scheduled_episodes.add(episode_index)
|
||||||
|
|
||||||
schedule_frontier()
|
def decode_item(
|
||||||
|
episode_future: Future[datasets.Dataset],
|
||||||
|
episode_index: int,
|
||||||
|
frame_index: int,
|
||||||
|
) -> dict:
|
||||||
|
return self._make_episode_item(
|
||||||
|
episode_future.result(),
|
||||||
|
episode_index,
|
||||||
|
frame_index,
|
||||||
|
video_cache=video_cache,
|
||||||
|
apply_image_transforms=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
def update_frontier() -> None:
|
||||||
|
for evicted_episode in planner.evicted:
|
||||||
|
episode_futures.pop(evicted_episode, None)
|
||||||
|
if video_cache is not None and evicted_episode in retained_video_episodes:
|
||||||
|
video_cache.release_episode(evicted_episode)
|
||||||
|
retained_video_episodes.remove(evicted_episode)
|
||||||
|
if video_cache is not None:
|
||||||
|
for admitted_episode in planner.newly_admitted:
|
||||||
|
if admitted_episode not in retained_video_episodes:
|
||||||
|
video_cache.retain_episode(admitted_episode)
|
||||||
|
retained_video_episodes.add(admitted_episode)
|
||||||
|
planner.evicted.clear()
|
||||||
|
planner.newly_admitted.clear()
|
||||||
|
schedule_frontier()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
while True:
|
schedule_frontier()
|
||||||
try:
|
planner_exhausted = False
|
||||||
episode_index, frame_index = next(planner)
|
while decoded_futures or not planner_exhausted:
|
||||||
except StopIteration:
|
while not planner_exhausted and len(decoded_futures) < self.decoded_queue_size:
|
||||||
self._active_epoch = epoch + 1 if self.shuffle else 0
|
try:
|
||||||
self._state_offset = 0
|
episode_index, frame_index = next(planner)
|
||||||
break
|
except StopIteration:
|
||||||
|
planner_exhausted = True
|
||||||
|
break
|
||||||
|
|
||||||
episode_dataset = submit(episode_index).result()
|
episode_future = submit(episode_index)
|
||||||
item = self._make_episode_item(
|
if video_cache is not None:
|
||||||
episode_dataset,
|
video_cache.retain_episode(episode_index)
|
||||||
episode_index,
|
try:
|
||||||
frame_index,
|
decoded_future = decode_executor.submit(
|
||||||
video_cache=video_cache,
|
decode_item,
|
||||||
)
|
episode_future,
|
||||||
|
episode_index,
|
||||||
for evicted_episode in planner.evicted:
|
frame_index,
|
||||||
episode_futures.pop(evicted_episode, None)
|
)
|
||||||
if video_cache is not None and evicted_episode in retained_video_episodes:
|
except Exception:
|
||||||
video_cache.release_episode(evicted_episode)
|
if video_cache is not None:
|
||||||
retained_video_episodes.remove(evicted_episode)
|
video_cache.release_episode(episode_index)
|
||||||
if video_cache is not None:
|
raise
|
||||||
for admitted_episode in planner.newly_admitted:
|
if video_cache is not None:
|
||||||
if admitted_episode not in retained_video_episodes:
|
decoded_future.add_done_callback(
|
||||||
video_cache.retain_episode(admitted_episode)
|
lambda _future, retained_episode=episode_index, retained_cache=video_cache: (
|
||||||
retained_video_episodes.add(admitted_episode)
|
retained_cache.release_episode(retained_episode)
|
||||||
planner.evicted.clear()
|
)
|
||||||
planner.newly_admitted.clear()
|
)
|
||||||
schedule_frontier()
|
decoded_futures.append(decoded_future)
|
||||||
|
update_frontier()
|
||||||
|
|
||||||
|
if not decoded_futures:
|
||||||
|
continue
|
||||||
|
item = decoded_futures.popleft().result()
|
||||||
|
self._apply_image_transforms(item)
|
||||||
self._state_offset += 1
|
self._state_offset += 1
|
||||||
yield item
|
yield item
|
||||||
|
self._active_epoch = epoch + 1 if self.shuffle else 0
|
||||||
|
self._state_offset = 0
|
||||||
finally:
|
finally:
|
||||||
|
for future in decoded_futures:
|
||||||
|
future.cancel()
|
||||||
|
decode_executor.shutdown(wait=True, cancel_futures=True)
|
||||||
for future in episode_futures.values():
|
for future in episode_futures.values():
|
||||||
future.cancel()
|
future.cancel()
|
||||||
executor.shutdown(wait=True, cancel_futures=True)
|
parquet_executor.shutdown(wait=True, cancel_futures=True)
|
||||||
if video_cache is not None:
|
if video_cache is not None:
|
||||||
video_cache.close()
|
video_cache.close()
|
||||||
|
|
||||||
@@ -483,14 +555,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
sidecar_path=self._sidecar_path,
|
sidecar_path=self._sidecar_path,
|
||||||
token=self._streaming_io_token,
|
token=self._streaming_io_token,
|
||||||
)
|
)
|
||||||
decoder_limit = max(1, min(64, self.episode_pool_size * max(1, len(self.meta.video_keys))))
|
|
||||||
return EpisodeByteCache(
|
return EpisodeByteCache(
|
||||||
manifest,
|
manifest,
|
||||||
self._data_root,
|
self._data_root,
|
||||||
byte_budget=self.byte_budget,
|
byte_budget=self.byte_budget,
|
||||||
workers=workers,
|
workers=workers,
|
||||||
range_backend=range_backend,
|
range_backend=range_backend,
|
||||||
max_open_decoders=decoder_limit,
|
native_http_connections=self.native_http_connections,
|
||||||
|
native_http_subranges=self.native_http_subranges,
|
||||||
|
max_open_decoders=self.max_open_decoders,
|
||||||
video_backend=self._video_backend,
|
video_backend=self._video_backend,
|
||||||
tolerance_s=self.tolerance_s,
|
tolerance_s=self.tolerance_s,
|
||||||
token=self._streaming_io_token,
|
token=self._streaming_io_token,
|
||||||
@@ -503,6 +576,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
frame_index: int,
|
frame_index: int,
|
||||||
*,
|
*,
|
||||||
video_cache: EpisodeByteCache | None,
|
video_cache: EpisodeByteCache | None,
|
||||||
|
apply_image_transforms: bool = True,
|
||||||
) -> dict:
|
) -> dict:
|
||||||
item = episode_dataset[frame_index]
|
item = episode_dataset[frame_index]
|
||||||
episode = self.meta.episodes[episode_index]
|
episode = self.meta.episodes[episode_index]
|
||||||
@@ -563,11 +637,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
frames = frames.to(torch.float32) / 255.0
|
frames = frames.to(torch.float32) / 255.0
|
||||||
item[video_key] = frames.squeeze(0)
|
item[video_key] = frames.squeeze(0)
|
||||||
|
|
||||||
if self.image_transforms is not None:
|
if apply_image_transforms:
|
||||||
for camera_key in self.meta.camera_keys:
|
self._apply_image_transforms(item)
|
||||||
if camera_key in self.meta.depth_keys:
|
|
||||||
continue
|
|
||||||
item[camera_key] = self.image_transforms(item[camera_key])
|
|
||||||
|
|
||||||
for key, stored_unit in self._image_depth_units.items():
|
for key, stored_unit in self._image_depth_units.items():
|
||||||
if key in item and stored_unit is not None and stored_unit != self._depth_output_unit:
|
if key in item and stored_unit is not None and stored_unit != self._depth_output_unit:
|
||||||
@@ -585,6 +656,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
)
|
)
|
||||||
return item
|
return item
|
||||||
|
|
||||||
|
def _apply_image_transforms(self, item: dict) -> None:
|
||||||
|
if self.image_transforms is None:
|
||||||
|
return
|
||||||
|
for camera_key in self.meta.camera_keys:
|
||||||
|
if camera_key in self.meta.depth_keys:
|
||||||
|
continue
|
||||||
|
item[camera_key] = self.image_transforms(item[camera_key])
|
||||||
|
|
||||||
def state_dict(self) -> dict[str, int]:
|
def state_dict(self) -> dict[str, int]:
|
||||||
return {
|
return {
|
||||||
"epoch": self._active_epoch,
|
"epoch": self._active_epoch,
|
||||||
|
|||||||
@@ -75,6 +75,7 @@ class EpisodeByteCache:
|
|||||||
self._pool = ThreadPoolExecutor(max_workers=workers)
|
self._pool = ThreadPoolExecutor(max_workers=workers)
|
||||||
self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict()
|
self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict()
|
||||||
self._decoders: OrderedDict[tuple[int, str], Any] = OrderedDict()
|
self._decoders: OrderedDict[tuple[int, str], Any] = OrderedDict()
|
||||||
|
self._decoder_locks: dict[tuple[int, str], threading.Lock] = {}
|
||||||
self._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {}
|
self._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {}
|
||||||
self._retained_episodes: dict[int, int] = {}
|
self._retained_episodes: dict[int, int] = {}
|
||||||
self._decoder_fallback_count = 0
|
self._decoder_fallback_count = 0
|
||||||
@@ -96,6 +97,7 @@ class EpisodeByteCache:
|
|||||||
decoders = list(self._decoders.values())
|
decoders = list(self._decoders.values())
|
||||||
self._cache.clear()
|
self._cache.clear()
|
||||||
self._decoders.clear()
|
self._decoders.clear()
|
||||||
|
self._decoder_locks.clear()
|
||||||
self._futures.clear()
|
self._futures.clear()
|
||||||
self._retained_episodes.clear()
|
self._retained_episodes.clear()
|
||||||
self._fallback_decoders.clear()
|
self._fallback_decoders.clear()
|
||||||
@@ -186,6 +188,7 @@ class EpisodeByteCache:
|
|||||||
self._decoders[key] = decoder
|
self._decoders[key] = decoder
|
||||||
while len(self._decoders) > self.max_open_decoders:
|
while len(self._decoders) > self.max_open_decoders:
|
||||||
evicted_key, evicted_decoder = self._decoders.popitem(last=False)
|
evicted_key, evicted_decoder = self._decoders.popitem(last=False)
|
||||||
|
self._decoder_locks.pop(evicted_key, None)
|
||||||
self._fallback_decoders.discard(evicted_key)
|
self._fallback_decoders.discard(evicted_key)
|
||||||
_close_decoder(evicted_decoder)
|
_close_decoder(evicted_decoder)
|
||||||
return decoder
|
return decoder
|
||||||
@@ -227,15 +230,23 @@ class EpisodeByteCache:
|
|||||||
decoder, release = self._decoder_for_frames(episode_index, camera_key)
|
decoder, release = self._decoder_for_frames(episode_index, camera_key)
|
||||||
with self._lock:
|
with self._lock:
|
||||||
uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders
|
uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders
|
||||||
|
decode_lock = self._decoder_locks.setdefault(key, threading.Lock())
|
||||||
try:
|
try:
|
||||||
if uses_pyav_timestamps:
|
with decode_lock:
|
||||||
return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data
|
if uses_pyav_timestamps:
|
||||||
metadata = decoder.metadata
|
return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data
|
||||||
fps = getattr(metadata, "average_fps", None)
|
metadata = decoder.metadata
|
||||||
if fps is None:
|
fps = getattr(metadata, "average_fps", None)
|
||||||
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
|
if fps is None:
|
||||||
fps = metadata.num_frames / duration
|
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
|
||||||
return decoder.get_frames_at(indices=[round(ts * fps) for ts in local_ts]).data
|
fps = metadata.num_frames / duration
|
||||||
|
num_frames = getattr(metadata, "num_frames", None)
|
||||||
|
if num_frames is None:
|
||||||
|
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
|
||||||
|
num_frames = round(duration * fps)
|
||||||
|
last_index = max(0, int(num_frames) - 1)
|
||||||
|
indices = [min(max(round(ts * fps), 0), last_index) for ts in local_ts]
|
||||||
|
return decoder.get_frames_at(indices=indices).data
|
||||||
finally:
|
finally:
|
||||||
if release is not None:
|
if release is not None:
|
||||||
release()
|
release()
|
||||||
@@ -257,6 +268,7 @@ class EpisodeByteCache:
|
|||||||
with self._lock:
|
with self._lock:
|
||||||
if self._decoders.get(key) is decoder:
|
if self._decoders.get(key) is decoder:
|
||||||
self._decoders.pop(key)
|
self._decoders.pop(key)
|
||||||
|
self._decoder_locks.pop(key, None)
|
||||||
self._fallback_decoders.discard(key)
|
self._fallback_decoders.discard(key)
|
||||||
continue
|
continue
|
||||||
return decoder, decoder.release
|
return decoder, decoder.release
|
||||||
@@ -330,6 +342,7 @@ class EpisodeByteCache:
|
|||||||
entry = self._cache.pop(key)
|
entry = self._cache.pop(key)
|
||||||
self._bytes -= len(entry["bytes"])
|
self._bytes -= len(entry["bytes"])
|
||||||
decoder = self._decoders.pop(key, None)
|
decoder = self._decoders.pop(key, None)
|
||||||
|
self._decoder_locks.pop(key, None)
|
||||||
self._fallback_decoders.discard(key)
|
self._fallback_decoders.discard(key)
|
||||||
if decoder is not None:
|
if decoder is not None:
|
||||||
_close_decoder(decoder)
|
_close_decoder(decoder)
|
||||||
|
|||||||
@@ -11,8 +11,11 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import posixpath
|
import posixpath
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
import fsspec
|
import fsspec
|
||||||
import pyarrow as pa
|
import pyarrow as pa
|
||||||
@@ -29,6 +32,8 @@ class EpisodeParquetReader:
|
|||||||
*,
|
*,
|
||||||
columns: Sequence[str],
|
columns: Sequence[str],
|
||||||
token: str | bool | None = None,
|
token: str | bool | None = None,
|
||||||
|
max_retries: int = 4,
|
||||||
|
retry_backoff_s: float = 0.05,
|
||||||
):
|
):
|
||||||
if not columns:
|
if not columns:
|
||||||
raise ValueError("EpisodeParquetReader requires at least one projected column")
|
raise ValueError("EpisodeParquetReader requires at least one projected column")
|
||||||
@@ -37,6 +42,13 @@ class EpisodeParquetReader:
|
|||||||
self.columns if "episode_index" in self.columns else (*self.columns, "episode_index")
|
self.columns if "episode_index" in self.columns else (*self.columns, "episode_index")
|
||||||
)
|
)
|
||||||
data_root_str = str(data_root)
|
data_root_str = str(data_root)
|
||||||
|
if max_retries < 0:
|
||||||
|
raise ValueError("max_retries must be non-negative")
|
||||||
|
if retry_backoff_s < 0:
|
||||||
|
raise ValueError("retry_backoff_s must be non-negative")
|
||||||
|
self._max_retries = max_retries
|
||||||
|
self._retry_backoff_s = retry_backoff_s
|
||||||
|
self._open_lock = threading.Lock() if data_root_str.startswith("hf://") else None
|
||||||
storage_options = {"token": token} if token is not None and data_root_str.startswith("hf://") else {}
|
storage_options = {"token": token} if token is not None and data_root_str.startswith("hf://") else {}
|
||||||
self._filesystem, self._root_path = fsspec.core.url_to_fs(data_root_str, **storage_options)
|
self._filesystem, self._root_path = fsspec.core.url_to_fs(data_root_str, **storage_options)
|
||||||
|
|
||||||
@@ -51,7 +63,7 @@ class EpisodeParquetReader:
|
|||||||
raise ValueError(f"Episode {episode_index} must contain at least one row")
|
raise ValueError(f"Episode {episode_index} must contain at least one row")
|
||||||
|
|
||||||
path = posixpath.join(self._root_path.rstrip("/"), str(relative_path).lstrip("/"))
|
path = posixpath.join(self._root_path.rstrip("/"), str(relative_path).lstrip("/"))
|
||||||
with self._filesystem.open(path, "rb") as source:
|
with self._open_with_retry(path) as source:
|
||||||
parquet = pq.ParquetFile(source)
|
parquet = pq.ParquetFile(source)
|
||||||
available = set(parquet.schema_arrow.names)
|
available = set(parquet.schema_arrow.names)
|
||||||
missing = sorted(set(self._read_columns) - available)
|
missing = sorted(set(self._read_columns) - available)
|
||||||
@@ -71,6 +83,20 @@ class EpisodeParquetReader:
|
|||||||
table = table.drop_columns(["episode_index"])
|
table = table.drop_columns(["episode_index"])
|
||||||
return table
|
return table
|
||||||
|
|
||||||
|
def _open_with_retry(self, path: str) -> Any:
|
||||||
|
for attempt in range(self._max_retries + 1):
|
||||||
|
try:
|
||||||
|
if self._open_lock is None:
|
||||||
|
return self._filesystem.open(path, "rb")
|
||||||
|
with self._open_lock:
|
||||||
|
return self._filesystem.open(path, "rb")
|
||||||
|
except FileNotFoundError:
|
||||||
|
if attempt == self._max_retries:
|
||||||
|
raise
|
||||||
|
self._filesystem.invalidate_cache(posixpath.dirname(path))
|
||||||
|
time.sleep(self._retry_backoff_s * 2**attempt)
|
||||||
|
raise RuntimeError("unreachable")
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None:
|
def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None:
|
||||||
episode_column = next(
|
episode_column = next(
|
||||||
|
|||||||
@@ -44,6 +44,11 @@ def test_dataset_config_empty_episodes_ok():
|
|||||||
("streaming_episode_pool_size", 0, "episode_pool_size"),
|
("streaming_episode_pool_size", 0, "episode_pool_size"),
|
||||||
("streaming_prefetch_episodes", -1, "prefetch_episodes"),
|
("streaming_prefetch_episodes", -1, "prefetch_episodes"),
|
||||||
("streaming_byte_budget_gb", 0, "byte_budget_gb"),
|
("streaming_byte_budget_gb", 0, "byte_budget_gb"),
|
||||||
|
("streaming_decode_threads", 0, "decode_threads"),
|
||||||
|
("streaming_decoded_queue_size", 0, "decoded_queue_size"),
|
||||||
|
("streaming_max_open_decoders", 0, "max_open_decoders"),
|
||||||
|
("streaming_native_http_connections", 0, "native_http_connections"),
|
||||||
|
("streaming_native_http_subranges", 0, "native_http_subranges"),
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message):
|
def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message):
|
||||||
|
|||||||
@@ -122,3 +122,33 @@ def test_reader_forwards_explicit_token_to_hf_filesystem(monkeypatch) -> None:
|
|||||||
)
|
)
|
||||||
|
|
||||||
url_to_fs.assert_called_once_with("hf://datasets/private@revision", token="hf_test_token")
|
url_to_fs.assert_called_once_with("hf://datasets/private@revision", token="hf_test_token")
|
||||||
|
|
||||||
|
|
||||||
|
def test_reader_retries_transient_remote_file_not_found(tmp_path: Path, monkeypatch) -> None:
|
||||||
|
path = tmp_path / "episode.parquet"
|
||||||
|
pq.write_table(_table([0, 0]), path)
|
||||||
|
filesystem = Mock()
|
||||||
|
attempts = 0
|
||||||
|
|
||||||
|
def open_remote(*_args, **_kwargs):
|
||||||
|
nonlocal attempts
|
||||||
|
attempts += 1
|
||||||
|
if attempts == 1:
|
||||||
|
raise FileNotFoundError("partial HfFileSystem dircache")
|
||||||
|
return path.open("rb")
|
||||||
|
|
||||||
|
filesystem.open.side_effect = open_remote
|
||||||
|
filesystem.invalidate_cache = Mock()
|
||||||
|
monkeypatch.setattr(fsspec.core, "url_to_fs", Mock(return_value=(filesystem, "root")))
|
||||||
|
reader = EpisodeParquetReader(
|
||||||
|
"hf://datasets/private@revision",
|
||||||
|
columns=("episode_index", "frame_index"),
|
||||||
|
max_retries=1,
|
||||||
|
retry_backoff_s=0,
|
||||||
|
)
|
||||||
|
|
||||||
|
table = reader.read_episode("data/file.parquet", episode_index=0, expected_rows=2)
|
||||||
|
|
||||||
|
assert len(table) == 2
|
||||||
|
assert attempts == 2
|
||||||
|
filesystem.invalidate_cache.assert_called_once()
|
||||||
|
|||||||
@@ -11,6 +11,7 @@
|
|||||||
import json
|
import json
|
||||||
import struct
|
import struct
|
||||||
import threading
|
import threading
|
||||||
|
import time
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -217,6 +218,66 @@ def test_decoder_falls_back_to_pyav_when_torchcodec_rejects_mini_mp4(monkeypatch
|
|||||||
assert cache.decoder_fallback_count == 1
|
assert cache.decoder_fallback_count == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_torchcodec_frame_indices_are_clamped_to_decoder_bounds(monkeypatch, tmp_path):
|
||||||
|
requested_indices = []
|
||||||
|
|
||||||
|
class FakeDecoder:
|
||||||
|
metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})()
|
||||||
|
|
||||||
|
def get_frames_at(self, *, indices):
|
||||||
|
requested_indices.extend(indices)
|
||||||
|
return type("Frames", (), {"data": indices})()
|
||||||
|
|
||||||
|
with _fake_cache(monkeypatch, tmp_path) as cache:
|
||||||
|
monkeypatch.setattr(
|
||||||
|
cache.manifest,
|
||||||
|
"lookup",
|
||||||
|
lambda *_args: type("Span", (), {"source_start_pts": 0.0})(),
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (FakeDecoder(), None))
|
||||||
|
|
||||||
|
cache.get_frames(0, "camera", [-0.1, 1.0])
|
||||||
|
|
||||||
|
assert requested_indices == [0, 9]
|
||||||
|
|
||||||
|
|
||||||
|
def test_frame_reads_serialize_access_to_each_decoder(monkeypatch, tmp_path):
|
||||||
|
state_lock = threading.Lock()
|
||||||
|
active = 0
|
||||||
|
max_active = 0
|
||||||
|
|
||||||
|
class FakeDecoder:
|
||||||
|
metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})()
|
||||||
|
|
||||||
|
def get_frames_at(self, *, indices):
|
||||||
|
nonlocal active, max_active
|
||||||
|
with state_lock:
|
||||||
|
active += 1
|
||||||
|
max_active = max(max_active, active)
|
||||||
|
try:
|
||||||
|
time.sleep(0.01)
|
||||||
|
return type("Frames", (), {"data": indices})()
|
||||||
|
finally:
|
||||||
|
with state_lock:
|
||||||
|
active -= 1
|
||||||
|
|
||||||
|
with _fake_cache(monkeypatch, tmp_path) as cache:
|
||||||
|
monkeypatch.setattr(
|
||||||
|
cache.manifest,
|
||||||
|
"lookup",
|
||||||
|
lambda *_args: type("Span", (), {"source_start_pts": 0.0})(),
|
||||||
|
)
|
||||||
|
decoder = FakeDecoder()
|
||||||
|
monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (decoder, None))
|
||||||
|
|
||||||
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||||
|
futures = [executor.submit(cache.get_frames, 0, "camera", [0.1]) for _ in range(2)]
|
||||||
|
for future in futures:
|
||||||
|
future.result()
|
||||||
|
|
||||||
|
assert max_active == 1
|
||||||
|
|
||||||
|
|
||||||
def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path):
|
def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path):
|
||||||
with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache:
|
with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache:
|
||||||
cache.retain_episode(0)
|
cache.retain_episode(0)
|
||||||
|
|||||||
@@ -38,6 +38,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
|
|||||||
streaming_episode_pool_size=7,
|
streaming_episode_pool_size=7,
|
||||||
streaming_prefetch_episodes=3,
|
streaming_prefetch_episodes=3,
|
||||||
streaming_byte_budget_gb=2.5,
|
streaming_byte_budget_gb=2.5,
|
||||||
|
streaming_decode_threads=2,
|
||||||
|
streaming_decoded_queue_size=5,
|
||||||
|
streaming_max_open_decoders=17,
|
||||||
|
streaming_native_http_connections=9,
|
||||||
|
streaming_native_http_subranges=3,
|
||||||
)
|
)
|
||||||
cfg = SimpleNamespace(
|
cfg = SimpleNamespace(
|
||||||
dataset=dataset_config,
|
dataset=dataset_config,
|
||||||
@@ -54,6 +59,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
|
|||||||
assert captured["kwargs"]["episode_pool_size"] == 7
|
assert captured["kwargs"]["episode_pool_size"] == 7
|
||||||
assert captured["kwargs"]["prefetch_episodes"] == 3
|
assert captured["kwargs"]["prefetch_episodes"] == 3
|
||||||
assert captured["kwargs"]["byte_budget_gb"] == 2.5
|
assert captured["kwargs"]["byte_budget_gb"] == 2.5
|
||||||
|
assert captured["kwargs"]["decode_threads"] == 2
|
||||||
|
assert captured["kwargs"]["decoded_queue_size"] == 5
|
||||||
|
assert captured["kwargs"]["max_open_decoders"] == 17
|
||||||
|
assert captured["kwargs"]["native_http_connections"] == 9
|
||||||
|
assert captured["kwargs"]["native_http_subranges"] == 3
|
||||||
assert captured["kwargs"]["max_num_shards"] == 1
|
assert captured["kwargs"]["max_num_shards"] == 1
|
||||||
assert captured["kwargs"]["video_backend"] == "pyav"
|
assert captured["kwargs"]["video_backend"] == "pyav"
|
||||||
assert captured["kwargs"]["return_uint8"] is True
|
assert captured["kwargs"]["return_uint8"] is True
|
||||||
|
|||||||
@@ -10,6 +10,8 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
from itertools import islice
|
from itertools import islice
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
@@ -61,6 +63,61 @@ def test_streaming_matches_map_style_with_exact_coverage(tmp_path: Path, lerobot
|
|||||||
_assert_item_equal(sample, map_dataset[int(sample["index"])])
|
_assert_item_equal(sample, map_dataset[int(sample["index"])])
|
||||||
|
|
||||||
|
|
||||||
|
def test_parallel_decode_queue_preserves_planner_order(
|
||||||
|
tmp_path: Path,
|
||||||
|
lerobot_dataset_factory,
|
||||||
|
monkeypatch,
|
||||||
|
) -> None:
|
||||||
|
root = tmp_path / "dataset"
|
||||||
|
lerobot_dataset_factory(
|
||||||
|
root=root,
|
||||||
|
repo_id=DUMMY_REPO_ID,
|
||||||
|
total_episodes=4,
|
||||||
|
total_frames=40,
|
||||||
|
use_videos=False,
|
||||||
|
)
|
||||||
|
expected = _indices(
|
||||||
|
StreamingLeRobotDataset(
|
||||||
|
DUMMY_REPO_ID,
|
||||||
|
root=root,
|
||||||
|
seed=17,
|
||||||
|
buffer_size=3,
|
||||||
|
decode_threads=1,
|
||||||
|
decoded_queue_size=1,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
parallel = StreamingLeRobotDataset(
|
||||||
|
DUMMY_REPO_ID,
|
||||||
|
root=root,
|
||||||
|
seed=17,
|
||||||
|
buffer_size=3,
|
||||||
|
decode_threads=3,
|
||||||
|
decoded_queue_size=5,
|
||||||
|
)
|
||||||
|
original_make_item = parallel._make_episode_item
|
||||||
|
state_lock = threading.Lock()
|
||||||
|
active = 0
|
||||||
|
max_active = 0
|
||||||
|
|
||||||
|
def delayed_make_item(*args, **kwargs):
|
||||||
|
nonlocal active, max_active
|
||||||
|
with state_lock:
|
||||||
|
active += 1
|
||||||
|
max_active = max(max_active, active)
|
||||||
|
try:
|
||||||
|
frame_index = int(args[2])
|
||||||
|
time.sleep(0.005 if frame_index % 3 == 0 else 0.001)
|
||||||
|
return original_make_item(*args, **kwargs)
|
||||||
|
finally:
|
||||||
|
with state_lock:
|
||||||
|
active -= 1
|
||||||
|
|
||||||
|
monkeypatch.setattr(parallel, "_make_episode_item", delayed_make_item)
|
||||||
|
|
||||||
|
assert _indices(parallel) == expected
|
||||||
|
assert 1 < max_active <= parallel.decode_threads
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
|
@pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
|
||||||
def test_streaming_rgb_video_matches_map_style(
|
def test_streaming_rgb_video_matches_map_style(
|
||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
|
|||||||
Reference in New Issue
Block a user