diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py index b9b9c6c43..e40a7c479 100644 --- a/src/lerobot/processor/pipeline.py +++ b/src/lerobot/processor/pipeline.py @@ -713,6 +713,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): ProcessorMigrationError: If the model requires migration to processor format. """ model_id = str(pretrained_model_name_or_path) + model_path = Path(model_id) + is_local_source = model_path.is_dir() or model_path.is_file() hub_download_kwargs = { "force_download": force_download, "resume_download": resume_download, @@ -731,7 +733,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): # 3. Build steps with overrides steps, validated_overrides = cls._build_steps_with_overrides( - loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs + loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source ) # 4. Validate that all overrides were used @@ -921,6 +923,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): model_id: str, base_path: Path | None, hub_download_kwargs: dict[str, Any], + is_local_source: bool = False, ) -> tuple[list[ProcessorStep], set[str]]: """Build all processor steps with overrides and state loading. @@ -944,7 +947,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): 3. **State Loading** (via _load_step_state): - **If step has "state_file"**: Load tensor state from .safetensors - **Local first**: Check base_path/state_file.safetensors - - **Hub fallback**: Download state file if not found locally + - **Hub fallback**: Download state file if the pipeline was loaded from the Hub - **Optional**: Only load if step has load_state_dict method 4. **Override Tracking**: @@ -962,6 +965,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): model_id: The model identifier (needed for Hub state file downloads) base_path: Local directory path for finding state files hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.) + is_local_source: Whether model_id resolved to a local directory or config file. Returns: Tuple of (instantiated_steps_list, unused_override_keys) @@ -975,7 +979,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides) for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True): - cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs) + cls._load_step_state( + step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source + ) return steps, remaining_override_keys @@ -1139,6 +1145,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): model_id: str, base_path: Path | None, hub_download_kwargs: dict[str, Any], + is_local_source: bool = False, ) -> None: """Load state dictionary for a processor step if available. @@ -1157,7 +1164,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): - **Use case**: Loading from local saved model directory 2. **Hub download fallback**: Download state file from repository - - **When triggered**: Local file not found or base_path is None + - **When triggered**: Local file not found and the pipeline source is a Hub repo - **Process**: Use hf_hub_download with same parameters as config - **Example**: Download "normalize_step_0.safetensors" from "user/repo" - **Result**: Downloaded to local cache, path returned @@ -1178,6 +1185,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): model_id: The model identifier (used for Hub downloads if needed) base_path: Local directory path for finding state files (None for Hub-only) hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.) + is_local_source: Whether model_id resolved to a local directory or config file. Note: This method modifies step_instance in-place and returns None. @@ -1191,6 +1199,12 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin): # Try local file first if base_path and (base_path / state_filename).exists(): state_path = str(base_path / state_filename) + elif is_local_source: + state_path = base_path / state_filename if base_path else Path(state_filename) + raise FileNotFoundError( + f"State file '{state_filename}' was not found for local processor pipeline " + f"'{model_id}' at '{state_path}'." + ) else: # Download from Hub state_path = hf_hub_download( diff --git a/tests/processor/test_pipeline_from_pretrained_helpers.py b/tests/processor/test_pipeline_from_pretrained_helpers.py index 89d45cbad..36f9a8fad 100644 --- a/tests/processor/test_pipeline_from_pretrained_helpers.py +++ b/tests/processor/test_pipeline_from_pretrained_helpers.py @@ -26,8 +26,17 @@ import tempfile from pathlib import Path import pytest +import torch +from safetensors.torch import save_file -from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError +from lerobot.configs import PipelineFeatureType, PolicyFeature +from lerobot.processor.pipeline import ( + DataProcessorPipeline, + ProcessorMigrationError, + ProcessorStep, + ProcessorStepRegistry, +) +from lerobot.types import EnvTransition # Simplified Config Loading Tests @@ -98,6 +107,140 @@ def test_load_config_nonexistent_path_tries_hub(): DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {}) +def test_from_pretrained_local_directory_missing_state_does_not_call_hub(monkeypatch): + """Local processor dirs must fail locally when a state file is missing.""" + + @ProcessorStepRegistry.register("local_missing_state_step") + class LocalMissingStateStep(ProcessorStep): + def __call__(self, transition: EnvTransition) -> EnvTransition: + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: + pass + + try: + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + config = { + "name": "LocalMissingStatePipeline", + "steps": [{"registry_name": "local_missing_state_step", "state_file": "missing.safetensors"}], + } + (tmp_path / "processor.json").write_text(json.dumps(config)) + + def fail_hub_download(*args, **kwargs): + pytest.fail("local missing processor state should not call hf_hub_download") + + monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download) + + with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"): + DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json") + finally: + ProcessorStepRegistry.unregister("local_missing_state_step") + + +def test_from_pretrained_local_config_file_missing_state_does_not_call_hub(monkeypatch): + """Local single-file processor configs must also keep missing state resolution local.""" + + @ProcessorStepRegistry.register("local_file_missing_state_step") + class LocalFileMissingStateStep(ProcessorStep): + def __call__(self, transition: EnvTransition) -> EnvTransition: + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: + pass + + try: + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_path = Path(tmp_dir) + config_path = tmp_path / "processor.json" + config = { + "name": "LocalFileMissingStatePipeline", + "steps": [ + {"registry_name": "local_file_missing_state_step", "state_file": "missing.safetensors"} + ], + } + config_path.write_text(json.dumps(config)) + + def fail_hub_download(*args, **kwargs): + pytest.fail("local missing processor state should not call hf_hub_download") + + monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download) + + with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"): + DataProcessorPipeline.from_pretrained(config_path, config_filename="ignored.json") + finally: + ProcessorStepRegistry.unregister("local_file_missing_state_step") + + +def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypatch, tmp_path): + """Hub sources still fall back to hf_hub_download for state files.""" + + @ProcessorStepRegistry.register("hub_state_step") + class HubStateStep(ProcessorStep): + def __init__(self): + self.value = torch.tensor(0) + + def __call__(self, transition: EnvTransition) -> EnvTransition: + return transition + + def transform_features( + self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]] + ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]: + return features + + def load_state_dict(self, state: dict[str, torch.Tensor]) -> None: + self.value = state["value"] + + try: + state_path = tmp_path / "downloaded.safetensors" + save_file({"value": torch.tensor(7)}, state_path) + loaded_config = { + "name": "HubStatePipeline", + "steps": [{"registry_name": "hub_state_step", "state_file": "hub_state.safetensors"}], + } + calls = [] + + def fake_load_config(cls, model_id, config_filename, hub_download_kwargs): + return loaded_config, tmp_path / "hub_cache" + + def fake_hub_download(**kwargs): + calls.append(kwargs) + return str(state_path) + + monkeypatch.setattr(DataProcessorPipeline, "_load_config", classmethod(fake_load_config)) + monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download) + + pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json") + + assert calls == [ + { + "repo_id": "user/repo", + "filename": "hub_state.safetensors", + "repo_type": "model", + "force_download": False, + "resume_download": None, + "proxies": None, + "token": None, + "cache_dir": None, + "local_files_only": False, + "revision": None, + } + ] + assert pipeline.steps[0].value.item() == 7 + finally: + ProcessorStepRegistry.unregister("hub_state_step") + + # Config Validation Tests