mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 12:15:59 +00:00
Derive streaming decoder cap from episode pool
This commit is contained in:
@@ -59,13 +59,13 @@ 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_decode_threads` | 2 | Parallel sample assembly and video decode workers |
|
| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers |
|
||||||
| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order |
|
| `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_max_open_decoders` | pool × cameras | Independent open-decoder LRU cap per rank |
|
||||||
| `streaming_native_http_connections` | unset | Native HTTP connection 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_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read |
|
||||||
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
||||||
@@ -84,7 +84,6 @@ lerobot-train \
|
|||||||
--dataset.streaming_byte_budget_gb=4 \
|
--dataset.streaming_byte_budget_gb=4 \
|
||||||
--dataset.streaming_decode_threads=2 \
|
--dataset.streaming_decode_threads=2 \
|
||||||
--dataset.streaming_decoded_queue_size=8 \
|
--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
|
||||||
|
|||||||
@@ -52,7 +52,7 @@ def parse_args() -> argparse.Namespace:
|
|||||||
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("--decode-threads", type=int, default=2)
|
||||||
parser.add_argument("--decoded-queue-size", type=int, default=8)
|
parser.add_argument("--decoded-queue-size", type=int, default=8)
|
||||||
parser.add_argument("--max-open-decoders", type=int, default=64)
|
parser.add_argument("--max-open-decoders", type=int, default=None)
|
||||||
parser.add_argument("--native-http-connections", type=int, default=None)
|
parser.add_argument("--native-http-connections", type=int, default=None)
|
||||||
parser.add_argument("--native-http-subranges", type=int, default=1)
|
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)
|
||||||
@@ -196,7 +196,8 @@ def main() -> None:
|
|||||||
"byte_budget_gb": args.byte_budget_gb,
|
"byte_budget_gb": args.byte_budget_gb,
|
||||||
"decode_threads": args.decode_threads,
|
"decode_threads": args.decode_threads,
|
||||||
"decoded_queue_size": args.decoded_queue_size,
|
"decoded_queue_size": args.decoded_queue_size,
|
||||||
"max_open_decoders": args.max_open_decoders,
|
"max_open_decoders": dataset.max_open_decoders,
|
||||||
|
"requested_max_open_decoders": args.max_open_decoders,
|
||||||
"native_http_connections": args.native_http_connections,
|
"native_http_connections": args.native_http_connections,
|
||||||
"native_http_subranges": args.native_http_subranges,
|
"native_http_subranges": args.native_http_subranges,
|
||||||
"warmup_batches": args.warmup_batches,
|
"warmup_batches": args.warmup_batches,
|
||||||
|
|||||||
@@ -55,8 +55,8 @@ class DatasetConfig:
|
|||||||
# Parallel sample assembly/decode workers and their bounded in-order result queue.
|
# Parallel sample assembly/decode workers and their bounded in-order result queue.
|
||||||
streaming_decode_threads: int = 2
|
streaming_decode_threads: int = 2
|
||||||
streaming_decoded_queue_size: int = 8
|
streaming_decoded_queue_size: int = 8
|
||||||
# Independent decoder-state cap.
|
# Independent decoder-state cap. None covers every camera in the configured episode pool.
|
||||||
streaming_max_open_decoders: int = 64
|
streaming_max_open_decoders: int | None = None
|
||||||
# Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
|
# Per-rank native HTTP limits. None preserves the fetcher's worker-derived default.
|
||||||
streaming_native_http_connections: int | None = None
|
streaming_native_http_connections: int | None = None
|
||||||
streaming_native_http_subranges: int = 1
|
streaming_native_http_subranges: int = 1
|
||||||
@@ -80,7 +80,7 @@ class DatasetConfig:
|
|||||||
raise ValueError("streaming_decode_threads must be positive")
|
raise ValueError("streaming_decode_threads must be positive")
|
||||||
if self.streaming_decoded_queue_size <= 0:
|
if self.streaming_decoded_queue_size <= 0:
|
||||||
raise ValueError("streaming_decoded_queue_size must be positive")
|
raise ValueError("streaming_decoded_queue_size must be positive")
|
||||||
if self.streaming_max_open_decoders <= 0:
|
if self.streaming_max_open_decoders is not None and self.streaming_max_open_decoders <= 0:
|
||||||
raise ValueError("streaming_max_open_decoders must be positive")
|
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:
|
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")
|
raise ValueError("streaming_native_http_connections must be positive")
|
||||||
|
|||||||
@@ -125,7 +125,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
token: str | bool | None = None,
|
token: str | bool | None = None,
|
||||||
decode_threads: int = 2,
|
decode_threads: int = 2,
|
||||||
decoded_queue_size: int = 8,
|
decoded_queue_size: int = 8,
|
||||||
max_open_decoders: int = 64,
|
max_open_decoders: int | None = None,
|
||||||
native_http_connections: int | None = None,
|
native_http_connections: int | None = None,
|
||||||
native_http_subranges: int = 1,
|
native_http_subranges: int = 1,
|
||||||
):
|
):
|
||||||
@@ -161,7 +161,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
decode_threads (int, optional): Parallel sample assembly and video decode workers.
|
decode_threads (int, optional): Parallel sample assembly and video decode workers.
|
||||||
decoded_queue_size (int, optional): Maximum number of decoded samples produced ahead
|
decoded_queue_size (int, optional): Maximum number of decoded samples produced ahead
|
||||||
of the consumer. Results are yielded in exact planner order.
|
of the consumer. Results are yielded in exact planner order.
|
||||||
max_open_decoders (int, optional): Maximum number of open video decoders per rank.
|
max_open_decoders (int | None, optional): Maximum number of open video decoders per
|
||||||
|
rank. By default every camera in the configured episode pool can remain open.
|
||||||
native_http_connections (int | None, optional): Native HTTP connection limit per rank.
|
native_http_connections (int | None, optional): Native HTTP connection limit per rank.
|
||||||
``None`` preserves the fetcher's worker-derived default.
|
``None`` preserves the fetcher's worker-derived default.
|
||||||
native_http_subranges (int, optional): Concurrent HTTP subranges per video range read.
|
native_http_subranges (int, optional): Concurrent HTTP subranges per video range read.
|
||||||
@@ -213,7 +214,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
raise ValueError("decode_threads must be positive")
|
raise ValueError("decode_threads must be positive")
|
||||||
if decoded_queue_size <= 0:
|
if decoded_queue_size <= 0:
|
||||||
raise ValueError("decoded_queue_size must be positive")
|
raise ValueError("decoded_queue_size must be positive")
|
||||||
if max_open_decoders <= 0:
|
if max_open_decoders is not None and max_open_decoders <= 0:
|
||||||
raise ValueError("max_open_decoders must be positive")
|
raise ValueError("max_open_decoders must be positive")
|
||||||
if native_http_connections is not None and native_http_connections <= 0:
|
if native_http_connections is not None and native_http_connections <= 0:
|
||||||
raise ValueError("native_http_connections must be positive")
|
raise ValueError("native_http_connections must be positive")
|
||||||
@@ -224,7 +225,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
self.byte_budget = int(byte_budget_gb * 1024**3)
|
self.byte_budget = int(byte_budget_gb * 1024**3)
|
||||||
self.decode_threads = decode_threads
|
self.decode_threads = decode_threads
|
||||||
self.decoded_queue_size = decoded_queue_size
|
self.decoded_queue_size = decoded_queue_size
|
||||||
self.max_open_decoders = max_open_decoders
|
|
||||||
self.native_http_connections = native_http_connections
|
self.native_http_connections = native_http_connections
|
||||||
self.native_http_subranges = native_http_subranges
|
self.native_http_subranges = native_http_subranges
|
||||||
self.repeat = repeat
|
self.repeat = repeat
|
||||||
@@ -250,6 +250,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
self.meta.rescale_depth_stats(self._depth_output_unit)
|
self.meta.rescale_depth_stats(self._depth_output_unit)
|
||||||
# Check version
|
# Check version
|
||||||
check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
|
check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
|
||||||
|
self.max_open_decoders = (
|
||||||
|
max_open_decoders
|
||||||
|
if max_open_decoders is not None
|
||||||
|
else max(1, self.episode_pool_size * len(self.meta.video_keys))
|
||||||
|
)
|
||||||
|
|
||||||
self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
|
self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
|
||||||
vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info"))
|
vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info"))
|
||||||
|
|||||||
@@ -38,6 +38,10 @@ def test_dataset_config_empty_episodes_ok():
|
|||||||
DatasetConfig(repo_id="user/repo", episodes=[])
|
DatasetConfig(repo_id="user/repo", episodes=[])
|
||||||
|
|
||||||
|
|
||||||
|
def test_dataset_config_derives_streaming_decoder_limit_by_default():
|
||||||
|
assert DatasetConfig(repo_id="user/repo").streaming_max_open_decoders is None
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
("field", "value", "message"),
|
("field", "value", "message"),
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -118,6 +118,35 @@ def test_parallel_decode_queue_preserves_planner_order(
|
|||||||
assert 1 < max_active <= parallel.decode_threads
|
assert 1 < max_active <= parallel.decode_threads
|
||||||
|
|
||||||
|
|
||||||
|
def test_default_decoder_limit_covers_the_configured_episode_pool(
|
||||||
|
tmp_path: Path,
|
||||||
|
lerobot_dataset_factory,
|
||||||
|
) -> None:
|
||||||
|
root = tmp_path / "dataset"
|
||||||
|
lerobot_dataset_factory(
|
||||||
|
root=root,
|
||||||
|
repo_id=DUMMY_REPO_ID,
|
||||||
|
total_episodes=2,
|
||||||
|
total_frames=20,
|
||||||
|
)
|
||||||
|
|
||||||
|
streaming = StreamingLeRobotDataset(
|
||||||
|
DUMMY_REPO_ID,
|
||||||
|
root=root,
|
||||||
|
episode_pool_size=7,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert streaming.max_open_decoders == 7 * len(streaming.meta.video_keys)
|
||||||
|
|
||||||
|
overridden = StreamingLeRobotDataset(
|
||||||
|
DUMMY_REPO_ID,
|
||||||
|
root=root,
|
||||||
|
episode_pool_size=7,
|
||||||
|
max_open_decoders=5,
|
||||||
|
)
|
||||||
|
assert overridden.max_open_decoders == 5
|
||||||
|
|
||||||
|
|
||||||
@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