From 0d383d09f2051444de211739196a28cc94736861 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Fri, 24 Jul 2026 18:51:35 +0200 Subject: [PATCH] feat(dataset): accept token argument for private HF Hub datasets (#4136) --- docs/source/using_dataset_tools.mdx | 4 +++ src/lerobot/datasets/dataset_metadata.py | 18 ++++++++-- src/lerobot/datasets/lerobot_dataset.py | 35 +++++++++++++++--- src/lerobot/datasets/multi_dataset.py | 3 ++ src/lerobot/datasets/streaming_dataset.py | 15 +++++++- src/lerobot/datasets/utils.py | 17 ++++++--- tests/datasets/test_dataset_utils.py | 31 +++++++++++++++- tests/datasets/test_lerobot_dataset.py | 43 +++++++++++++++++++++++ tests/datasets/test_streaming.py | 38 ++++++++++++++++++++ 9 files changed, 191 insertions(+), 13 deletions(-) diff --git a/docs/source/using_dataset_tools.mdx b/docs/source/using_dataset_tools.mdx index a6dcdb1a7..3ddc320a5 100644 --- a/docs/source/using_dataset_tools.mdx +++ b/docs/source/using_dataset_tools.mdx @@ -252,6 +252,10 @@ lerobot-dataset-viz \ --episode-index 0 ``` +For a private or gated dataset, authenticate first with `hf auth login`, or set the +`HF_TOKEN` environment variable. The Hub client then discovers the credential +automatically; no token argument is needed. + **From a local folder:** Add the `--root` option and set `--mode local`. For example, to search in `./my_local_data_dir/lerobot/pusht`: diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index 6e19d14fb..ed5e13833 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/src/lerobot/datasets/dataset_metadata.py @@ -73,6 +73,8 @@ class LeRobotDatasetMetadata: revision: str | None = None, force_cache_sync: bool = False, metadata_buffer_size: int = 10, + *, + token: str | bool | None = None, ): """Load or download metadata for an existing LeRobot dataset. @@ -94,6 +96,10 @@ class LeRobotDatasetMetadata: even when local files exist. metadata_buffer_size: Number of episode metadata records to buffer in memory before flushing to parquet. + token: Authentication token used for Hub requests. Pass a string + token, ``True`` to require the locally stored token, ``False`` + to disable authentication, or ``None`` to use the Hugging Face + Hub default. """ self.repo_id = repo_id self.revision = revision if revision else CODEBASE_VERSION @@ -113,9 +119,12 @@ class LeRobotDatasetMetadata: self._load_metadata() except (FileNotFoundError, NotADirectoryError): if is_valid_version(self.revision): - self.revision = get_safe_version(self.repo_id, self.revision) + if token is None: + self.revision = get_safe_version(self.repo_id, self.revision) + else: + self.revision = get_safe_version(self.repo_id, self.revision, token=token) - self._pull_from_repo(allow_patterns="meta/") + self._pull_from_repo(allow_patterns="meta/", token=token) self._load_metadata() def _flush_metadata_buffer(self) -> None: @@ -220,7 +229,10 @@ class LeRobotDatasetMetadata: self, allow_patterns: list[str] | str | None = None, ignore_patterns: list[str] | str | None = None, + *, + token: str | bool | None = None, ) -> None: + token_kwargs = {} if token is None else {"token": token} if self._requested_root is None: self.root = Path( snapshot_download( @@ -230,6 +242,7 @@ class LeRobotDatasetMetadata: cache_dir=HF_LEROBOT_HUB_CACHE, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, + **token_kwargs, ) ) return @@ -242,6 +255,7 @@ class LeRobotDatasetMetadata: local_dir=self._requested_root, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, + **token_kwargs, ) self.root = self._requested_root diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 77b3032df..4c44bc543 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -65,6 +65,8 @@ class LeRobotDataset(torch.utils.data.Dataset): encoder_threads: int | None = None, streaming_encoding: bool = False, encoder_queue_maxsize: int = 30, + *, + token: str | bool | None = None, ): """ 2 modes are available for instantiating this class, depending on 2 different use cases: @@ -197,6 +199,11 @@ class LeRobotDataset(torch.utils.data.Dataset): instead of writing PNG images first. This makes save_episode() near-instant. Defaults to False. encoder_queue_maxsize (int, optional): Maximum number of frames to buffer per camera when using streaming encoding. Defaults to 30 (~1s at 30fps). + token: Authentication token used while downloading this dataset + from the Hub. Pass a string token, ``True`` to require the + locally stored token, ``False`` to disable authentication, or + ``None`` to use the Hugging Face Hub default. The token is not + retained on the dataset instance after initialization. Note: Write-mode parameters (``streaming_encoding``, ``batch_encoding_size``) passed to @@ -220,7 +227,11 @@ class LeRobotDataset(torch.utils.data.Dataset): # Load metadata (sets self.root once from the resolved metadata root) self.meta = LeRobotDatasetMetadata( - self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync + self.repo_id, + self._requested_root, + self.revision, + force_cache_sync=force_cache_sync, + token=token, ) self.root = self.meta.root self.revision = self.meta.revision @@ -260,8 +271,11 @@ class LeRobotDataset(torch.utils.data.Dataset): # Load actual data if force_cache_sync or not self.reader.try_load(): if is_valid_version(self.revision): - self.revision = get_safe_version(self.repo_id, self.revision) - self._download(download_videos) + if token is None: + self.revision = get_safe_version(self.repo_id, self.revision) + else: + self.revision = get_safe_version(self.repo_id, self.revision, token=token) + self._download(download_videos, token=token) self.reader.load_and_activate() # Detect write-mode params for backward compatibility @@ -626,10 +640,11 @@ class LeRobotDataset(torch.utils.data.Dataset): hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset") hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset") - def _download(self, download_videos: bool = True) -> None: + def _download(self, download_videos: bool = True, *, token: str | bool | None = None) -> None: """Downloads the dataset from the given 'repo_id' at the provided version.""" ignore_patterns = None if download_videos else "videos/" files = None + token_kwargs = {} if token is None else {"token": token} if self.episodes is not None: # Reader is guaranteed to exist here (created in __init__ before _download) files = self.reader.get_episodes_file_paths() @@ -643,6 +658,7 @@ class LeRobotDataset(torch.utils.data.Dataset): cache_dir=HF_LEROBOT_HUB_CACHE, allow_patterns=files, ignore_patterns=ignore_patterns, + **token_kwargs, ) ) else: @@ -654,6 +670,7 @@ class LeRobotDataset(torch.utils.data.Dataset): local_dir=self._requested_root, allow_patterns=files, ignore_patterns=ignore_patterns, + **token_kwargs, ) self.meta.root = self._requested_root @@ -793,6 +810,8 @@ class LeRobotDataset(torch.utils.data.Dataset): image_writer_threads: int = 0, streaming_encoding: bool = False, encoder_queue_maxsize: int = 30, + *, + token: str | bool | None = None, ) -> "LeRobotDataset": """Resume recording on an existing dataset. @@ -826,6 +845,8 @@ class LeRobotDataset(torch.utils.data.Dataset): streaming_encoding: If ``True``, encode video in real-time during capture. encoder_queue_maxsize: Max buffered frames per camera for streaming. + token: Authentication token used if metadata must be downloaded + from the Hub. The token is not retained on the dataset instance. Returns: A :class:`LeRobotDataset` in write mode, ready to append episodes. @@ -854,7 +875,11 @@ class LeRobotDataset(torch.utils.data.Dataset): # Load metadata (revision-safe when root is not provided) obj.meta = LeRobotDatasetMetadata( - obj.repo_id, obj._requested_root, obj.revision, force_cache_sync=force_cache_sync + obj.repo_id, + obj._requested_root, + obj.revision, + force_cache_sync=force_cache_sync, + token=token, ) obj._encoder_threads = encoder_threads diff --git a/src/lerobot/datasets/multi_dataset.py b/src/lerobot/datasets/multi_dataset.py index b955c1114..cc15fdec7 100644 --- a/src/lerobot/datasets/multi_dataset.py +++ b/src/lerobot/datasets/multi_dataset.py @@ -48,6 +48,8 @@ class MultiLeRobotDataset(torch.utils.data.Dataset): tolerances_s: dict | None = None, download_videos: bool = True, video_backend: str | None = None, + *, + token: str | bool | None = None, ): super().__init__() self.repo_ids = repo_ids @@ -65,6 +67,7 @@ class MultiLeRobotDataset(torch.utils.data.Dataset): tolerance_s=self.tolerances_s[repo_id], download_videos=download_videos, video_backend=video_backend, + token=token, ) for repo_id in repo_ids ] diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 14d4a52a4..806f2c24c 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -256,6 +256,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): shuffle: bool = True, return_uint8: bool = False, depth_output_unit: str = DEFAULT_DEPTH_UNIT, + *, + token: str | bool | None = None, ): """Initialize a StreamingLeRobotDataset. @@ -278,6 +280,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True. depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm"). Defaults to "mm". + token: Authentication token used while streaming this dataset from + the Hub. Pass a string token, ``True`` to require the locally + stored token, ``False`` to disable authentication, or ``None`` + to use the Hugging Face Hub default. The token is not retained + on the dataset instance after initialization. """ super().__init__() self.repo_id = repo_id @@ -306,7 +313,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): # Load metadata self.meta = LeRobotDatasetMetadata( - self.repo_id, self._requested_root, self.revision, force_cache_sync=force_cache_sync + self.repo_id, + self._requested_root, + self.revision, + force_cache_sync=force_cache_sync, + token=token, ) self.root = self.meta.root self.revision = self.meta.revision @@ -334,12 +345,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.delta_timestamps = delta_timestamps self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps) + token_kwargs = {} if token is None or self.streaming_from_local else {"token": token} self.hf_dataset: datasets.IterableDataset = load_dataset( self.repo_id if not self.streaming_from_local else str(self.root), split="train", streaming=self.streaming, data_files="data/*/*.parquet", revision=self.revision, + **token_kwargs, ) self.num_shards = min(self.hf_dataset.num_shards, max_num_shards) diff --git a/src/lerobot/datasets/utils.py b/src/lerobot/datasets/utils.py index d30761515..bb31296ec 100644 --- a/src/lerobot/datasets/utils.py +++ b/src/lerobot/datasets/utils.py @@ -325,16 +325,19 @@ def check_version_compatibility( logging.warning(FUTURE_MESSAGE.format(repo_id=repo_id, version=v_check)) -def get_repo_versions(repo_id: str) -> list[packaging.version.Version]: +def get_repo_versions(repo_id: str, *, token: str | bool | None = None) -> list[packaging.version.Version]: """Return available valid versions (branches and tags) on a given Hub repo. Args: repo_id (str): The repository ID on the Hugging Face Hub. + token: Authentication token used for Hub requests. Pass a string token, + ``True`` to require the locally stored token, ``False`` to disable + authentication, or ``None`` to use the Hugging Face Hub default. Returns: list[packaging.version.Version]: A list of valid versions found. """ - api = HfApi() + api = HfApi() if token is None else HfApi(token=token) repo_refs = api.list_repo_refs(repo_id, repo_type="dataset") repo_refs = [b.name for b in repo_refs.branches + repo_refs.tags] repo_versions = [] @@ -345,7 +348,12 @@ def get_repo_versions(repo_id: str) -> list[packaging.version.Version]: return repo_versions -def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> str: +def get_safe_version( + repo_id: str, + version: str | packaging.version.Version, + *, + token: str | bool | None = None, +) -> str: """Return the specified version if available on repo, or the latest compatible one. If the exact version is not found, it looks for the latest version with the @@ -354,6 +362,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> Args: repo_id (str): The repository ID on the Hugging Face Hub. version (str | packaging.version.Version): The target version. + token: Authentication token forwarded to the Hub version lookup. Returns: str: The safe version string (e.g., "v1.2.3") to use as a revision. @@ -366,7 +375,7 @@ def get_safe_version(repo_id: str, version: str | packaging.version.Version) -> target_version = ( packaging.version.parse(version) if not isinstance(version, packaging.version.Version) else version ) - hub_versions = get_repo_versions(repo_id) + hub_versions = get_repo_versions(repo_id) if token is None else get_repo_versions(repo_id, token=token) if not hub_versions: raise RevisionNotFoundError( diff --git a/tests/datasets/test_dataset_utils.py b/tests/datasets/test_dataset_utils.py index bf705ba81..09d5af9aa 100644 --- a/tests/datasets/test_dataset_utils.py +++ b/tests/datasets/test_dataset_utils.py @@ -14,16 +14,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace +from unittest.mock import Mock + import pytest import torch +from packaging.version import Version pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") from datasets import Dataset # noqa: E402 from huggingface_hub import DatasetCard +import lerobot.datasets.utils as dataset_utils from lerobot.datasets.io_utils import hf_transform_to_torch -from lerobot.datasets.utils import create_lerobot_dataset_card +from lerobot.datasets.utils import create_lerobot_dataset_card, get_repo_versions, get_safe_version from lerobot.utils.constants import ACTION, OBS_IMAGES from lerobot.utils.feature_utils import combine_feature_dicts @@ -57,6 +62,30 @@ def test_default_parameters(): ] +@pytest.mark.parametrize("token", ["hf_test_token", True, False]) +def test_get_repo_versions_forwards_token(monkeypatch, token): + api = Mock() + api.list_repo_refs.return_value = SimpleNamespace( + branches=[SimpleNamespace(name="v3.0")], + tags=[], + ) + hf_api = Mock(return_value=api) + monkeypatch.setattr(dataset_utils, "HfApi", hf_api) + + assert get_repo_versions("private/repo", token=token) == [Version("3.0")] + hf_api.assert_called_once_with(token=token) + api.list_repo_refs.assert_called_once_with("private/repo", repo_type="dataset") + + +@pytest.mark.parametrize("token", ["hf_test_token", True, False]) +def test_get_safe_version_forwards_token(monkeypatch, token): + get_versions = Mock(return_value=[Version("3.0")]) + monkeypatch.setattr(dataset_utils, "get_repo_versions", get_versions) + + assert get_safe_version("private/repo", "v3.0", token=token) == "v3.0" + get_versions.assert_called_once_with("private/repo", token=token) + + def test_with_tags(): tags = ["tag1", "tag2"] card = create_lerobot_dataset_card(tags=tags) diff --git a/tests/datasets/test_lerobot_dataset.py b/tests/datasets/test_lerobot_dataset.py index f3bda037f..f1614de52 100644 --- a/tests/datasets/test_lerobot_dataset.py +++ b/tests/datasets/test_lerobot_dataset.py @@ -20,6 +20,7 @@ property delegation, and the full create-record-finalize-read lifecycle. """ from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -191,6 +192,48 @@ def test_metadata_without_root_uses_hub_cache_snapshot_download( } +@pytest.mark.parametrize("token", ["hf_test_token", True, False]) +def test_metadata_download_forwards_token(tmp_path, monkeypatch, token): + snapshot_root = tmp_path / "snapshot" + snapshot_download = Mock(return_value=str(snapshot_root)) + get_safe_version = Mock(return_value="v3.0") + load_metadata = Mock(side_effect=[FileNotFoundError, None]) + monkeypatch.setattr(dataset_metadata_module, "snapshot_download", snapshot_download) + monkeypatch.setattr(dataset_metadata_module, "get_safe_version", get_safe_version) + monkeypatch.setattr(LeRobotDatasetMetadata, "_load_metadata", load_metadata) + + meta = LeRobotDatasetMetadata( + repo_id=DUMMY_REPO_ID, + revision="v3.0", + token=token, + ) + + assert meta.root == snapshot_root + assert not hasattr(meta, "_token") + get_safe_version.assert_called_once_with(DUMMY_REPO_ID, "v3.0", token=token) + assert snapshot_download.call_args.kwargs["token"] is token + + +@pytest.mark.parametrize("token", ["hf_test_token", True, False]) +def test_data_download_forwards_token(tmp_path, monkeypatch, token): + snapshot_root = tmp_path / "snapshot" + snapshot_download = Mock(return_value=str(snapshot_root)) + monkeypatch.setattr(lerobot_dataset_module, "snapshot_download", snapshot_download) + + dataset = LeRobotDataset.__new__(LeRobotDataset) + dataset.repo_id = DUMMY_REPO_ID + dataset.revision = "main" + dataset.episodes = None + dataset._requested_root = None + dataset.meta = SimpleNamespace(root=None) + dataset.reader = SimpleNamespace(root=None) + + dataset._download(token=token) + + assert dataset.root == snapshot_root + assert snapshot_download.call_args.kwargs["token"] is token + + def test_without_root_reads_different_revisions_from_distinct_snapshot_roots( tmp_path, info_factory, diff --git a/tests/datasets/test_streaming.py b/tests/datasets/test_streaming.py index db167f657..cae4be5b6 100644 --- a/tests/datasets/test_streaming.py +++ b/tests/datasets/test_streaming.py @@ -13,12 +13,16 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace +from unittest.mock import Mock + import numpy as np import pytest import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") +import lerobot.datasets.streaming_dataset as streaming_dataset_module from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset from lerobot.datasets.utils import safe_shard from lerobot.utils.constants import ACTION @@ -71,6 +75,40 @@ def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int return expected_indices +@pytest.mark.parametrize("token", ["hf_test_token", True, False]) +@pytest.mark.parametrize("from_local", [False, True]) +def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, monkeypatch, token, from_local): + requested_root = tmp_path / "local" if from_local else None + metadata = SimpleNamespace( + root=requested_root or tmp_path / "snapshot", + revision=streaming_dataset_module.CODEBASE_VERSION, + _version=streaming_dataset_module.CODEBASE_VERSION, + features={}, + depth_keys=[], + image_keys=[], + rescale_depth_stats=Mock(), + ) + metadata_cls = Mock(return_value=metadata) + load_dataset = Mock(return_value=SimpleNamespace(num_shards=1)) + monkeypatch.setattr(streaming_dataset_module, "LeRobotDatasetMetadata", metadata_cls) + monkeypatch.setattr(streaming_dataset_module, "load_dataset", load_dataset) + + dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, root=requested_root, token=token) + + metadata_cls.assert_called_once_with( + DUMMY_REPO_ID, + requested_root, + streaming_dataset_module.CODEBASE_VERSION, + force_cache_sync=False, + token=token, + ) + if from_local: + assert "token" not in load_dataset.call_args.kwargs + else: + assert load_dataset.call_args.kwargs["token"] is token + assert not hasattr(dataset, "_token") + + def test_single_frame_consistency(tmp_path, lerobot_dataset_factory): """Test if are correctly accessed""" ds_num_frames = 400