chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,342 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Tests for processor migration detection functionality.
|
||||
"""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
|
||||
|
||||
def test_is_processor_config_valid_configs():
|
||||
"""Test processor config detection with valid configurations."""
|
||||
valid_configs = [
|
||||
{"steps": []}, # Empty steps
|
||||
{"steps": [{"class": "MyClass"}]}, # Class-based step
|
||||
{"steps": [{"registry_name": "my_step"}]}, # Registry-based step
|
||||
{"steps": [{"class": "A"}, {"registry_name": "B"}]}, # Mixed
|
||||
{"name": "Test", "steps": [{"class": "MyClass"}]}, # With name
|
||||
]
|
||||
|
||||
for i, config in enumerate(valid_configs):
|
||||
assert DataProcessorPipeline._is_processor_config(config), (
|
||||
f"Valid config {i} should be detected as processor config: {config}"
|
||||
)
|
||||
|
||||
|
||||
def test_is_processor_config_invalid_configs():
|
||||
"""Test processor config detection with invalid configurations."""
|
||||
invalid_configs = [
|
||||
{}, # No steps field
|
||||
{"steps": "not a list"}, # Steps is not a list
|
||||
{"steps": [{}]}, # Step without class or registry_name
|
||||
{"steps": ["not a dict"]}, # Step is not a dict
|
||||
{"steps": [{"other_field": "value"}]}, # Step with wrong fields
|
||||
{"other_field": "value"}, # Completely different structure
|
||||
]
|
||||
|
||||
for i, config in enumerate(invalid_configs):
|
||||
assert not DataProcessorPipeline._is_processor_config(config), (
|
||||
f"Invalid config {i} should not be detected as processor config: {config}"
|
||||
)
|
||||
|
||||
|
||||
def test_should_suggest_migration_with_processor_config():
|
||||
"""Test that migration is NOT suggested when processor config exists."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a valid processor config
|
||||
processor_config = {
|
||||
"name": "TestProcessor",
|
||||
"steps": [
|
||||
{
|
||||
"class": "lerobot.processor.normalize.NormalizeStep",
|
||||
"config": {"mean": 0.0, "std": 1.0},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
with open(tmp_path / "processor.json", "w") as f:
|
||||
json.dump(processor_config, f)
|
||||
|
||||
# Should NOT suggest migration (processor config exists)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_with_empty_processor_config():
|
||||
"""Test that migration is NOT suggested when empty processor config exists."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create an empty processor config
|
||||
empty_processor_config = {
|
||||
"name": "EmptyProcessor",
|
||||
"steps": [], # Empty steps is valid
|
||||
}
|
||||
|
||||
with open(tmp_path / "empty_processor.json", "w") as f:
|
||||
json.dump(empty_processor_config, f)
|
||||
|
||||
# Should NOT suggest migration (processor config exists, even if empty)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_with_model_config_only():
|
||||
"""Test that migration IS suggested when only model config exists."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a model config (like old LeRobot format)
|
||||
model_config = {
|
||||
"type": "act",
|
||||
"input_features": {OBS_STATE: {"shape": [7]}},
|
||||
"output_features": {ACTION: {"shape": [7]}},
|
||||
"hidden_dim": 256,
|
||||
"n_obs_steps": 1,
|
||||
"n_action_steps": 1,
|
||||
}
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# SHOULD suggest migration (model config exists but no processor)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert result
|
||||
|
||||
|
||||
def test_should_suggest_migration_no_json_files():
|
||||
"""Test that migration is NOT suggested when no JSON files exist."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create some non-JSON files
|
||||
with open(tmp_path / "model.safetensors", "w") as f:
|
||||
f.write("fake model data")
|
||||
|
||||
with open(tmp_path / "README.md", "w") as f:
|
||||
f.write("# Model README")
|
||||
|
||||
# Should NOT suggest migration (no JSON files)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_random_json_files():
|
||||
"""Test that migration IS suggested when JSON files exist but none are processor configs."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create some random JSON file (not a processor config)
|
||||
random_config = {"some_field": "some_value", "another_field": 123}
|
||||
|
||||
with open(tmp_path / "random.json", "w") as f:
|
||||
json.dump(random_config, f)
|
||||
|
||||
# SHOULD suggest migration (JSON files exist but none are processor configs)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert result
|
||||
|
||||
|
||||
def test_should_suggest_migration_mixed_configs():
|
||||
"""Test that migration is NOT suggested when processor config exists alongside other configs."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create both a processor config and a model config
|
||||
processor_config = {"name": "TestProcessor", "steps": [{"registry_name": "normalize_step"}]}
|
||||
|
||||
model_config = {"type": "diffusion", "hidden_dim": 512}
|
||||
|
||||
with open(tmp_path / "processor.json", "w") as f:
|
||||
json.dump(processor_config, f)
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# Should NOT suggest migration (processor config exists)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert not result
|
||||
|
||||
|
||||
def test_should_suggest_migration_invalid_json():
|
||||
"""Test that invalid JSON is handled gracefully."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create an invalid JSON file
|
||||
with open(tmp_path / "invalid.json", "w") as f:
|
||||
f.write("{ invalid json")
|
||||
|
||||
# Create a valid non-processor config
|
||||
model_config = {"type": "act"}
|
||||
with open(tmp_path / "model.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# SHOULD suggest migration (invalid JSON is ignored, but we have non-processor JSON)
|
||||
result = DataProcessorPipeline._should_suggest_migration(tmp_path)
|
||||
assert result
|
||||
|
||||
|
||||
def test_from_pretrained_multiple_json_files_migration_error():
|
||||
"""Test that multiple JSON files trigger ProcessorMigrationError."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create multiple non-processor configs
|
||||
model_config = {"type": "act", "hidden_dim": 128}
|
||||
train_config = {"batch_size": 32, "lr": 0.001}
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
with open(tmp_path / "train_config.json", "w") as f:
|
||||
json.dump(train_config, f)
|
||||
|
||||
# Should raise ProcessorMigrationError
|
||||
with pytest.raises(ProcessorMigrationError) as exc_info:
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
|
||||
|
||||
# Check the error details
|
||||
error = exc_info.value
|
||||
assert str(tmp_path) in str(error.model_path)
|
||||
assert "migrate_policy_normalization.py" in error.migration_command
|
||||
assert "not a valid processor configuration" in error.original_error
|
||||
|
||||
|
||||
def test_from_pretrained_no_processor_config_migration_error():
|
||||
"""Test that missing processor config triggers ProcessorMigrationError."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a model config but no processor
|
||||
model_config = {"type": "diffusion", "hidden_dim": 256}
|
||||
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# Should raise ProcessorMigrationError
|
||||
with pytest.raises(ProcessorMigrationError) as exc_info:
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
|
||||
|
||||
# Check the error details
|
||||
error = exc_info.value
|
||||
assert str(tmp_path) in str(error.model_path)
|
||||
assert "migrate_policy_normalization.py" in error.migration_command
|
||||
assert "not a valid processor configuration" in error.original_error
|
||||
|
||||
|
||||
def test_from_pretrained_valid_processor_no_migration_error():
|
||||
"""Test that valid processor config does NOT trigger migration error."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a valid processor config
|
||||
processor_config = {
|
||||
"name": "TestProcessor",
|
||||
"steps": [], # Empty is valid
|
||||
}
|
||||
|
||||
with open(tmp_path / "processor.json", "w") as f:
|
||||
json.dump(processor_config, f)
|
||||
|
||||
# Should succeed and create pipeline
|
||||
pipeline = DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
|
||||
assert pipeline is not None
|
||||
assert pipeline.name == "TestProcessor"
|
||||
assert len(pipeline) == 0
|
||||
|
||||
|
||||
def test_from_pretrained_no_json_files_no_migration_error():
|
||||
"""Test that directories with no JSON files don't trigger migration errors."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create some non-JSON files
|
||||
with open(tmp_path / "model.safetensors", "w") as f:
|
||||
f.write("fake model data")
|
||||
|
||||
# Should raise FileNotFoundError (config file not found)
|
||||
with pytest.raises(FileNotFoundError, match="not found in directory"):
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
|
||||
|
||||
|
||||
def test_processor_migration_error_creation():
|
||||
"""Test that ProcessorMigrationError is created correctly."""
|
||||
model_path = "/path/to/model"
|
||||
migration_command = "python migrate.py --path /path/to/model"
|
||||
original_error = "Config not found"
|
||||
|
||||
error = ProcessorMigrationError(model_path, migration_command, original_error)
|
||||
|
||||
assert error.model_path == model_path
|
||||
assert error.migration_command == migration_command
|
||||
assert error.original_error == original_error
|
||||
assert model_path in str(error)
|
||||
assert migration_command in str(error)
|
||||
assert original_error in str(error)
|
||||
|
||||
|
||||
def test_processor_migration_error_attributes():
|
||||
"""Test that ProcessorMigrationError has correct attributes."""
|
||||
model_path = Path("/test/path")
|
||||
migration_command = "python test.py"
|
||||
original_error = "Test error"
|
||||
|
||||
error = ProcessorMigrationError(model_path, migration_command, original_error)
|
||||
|
||||
# Test that attributes are accessible
|
||||
assert hasattr(error, "model_path")
|
||||
assert hasattr(error, "migration_command")
|
||||
assert hasattr(error, "original_error")
|
||||
|
||||
# Test that it's still an Exception
|
||||
assert isinstance(error, Exception)
|
||||
|
||||
|
||||
def test_migration_suggestion_raises_error():
|
||||
"""Test that migration suggestion always raises ProcessorMigrationError."""
|
||||
with pytest.raises(ProcessorMigrationError) as exc_info:
|
||||
DataProcessorPipeline._suggest_processor_migration("/test/path", "Test error")
|
||||
|
||||
error = exc_info.value
|
||||
assert "/test/path" in str(error.model_path)
|
||||
assert "Test error" in error.original_error
|
||||
assert "migrate_policy_normalization.py" in error.migration_command
|
||||
|
||||
|
||||
def test_migration_error_always_raised_for_invalid_configs():
|
||||
"""Test that ProcessorMigrationError is always raised for invalid configs."""
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
tmp_path = Path(tmp_dir)
|
||||
|
||||
# Create a model config
|
||||
model_config = {"type": "test", "param": "value"}
|
||||
with open(tmp_path / "config.json", "w") as f:
|
||||
json.dump(model_config, f)
|
||||
|
||||
# Should always raise ProcessorMigrationError
|
||||
with pytest.raises(ProcessorMigrationError):
|
||||
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="config.json")
|
||||
Reference in New Issue
Block a user