feat(dataset): accept token argument for private HF Hub datasets (#4136)

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)

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,

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