Fix: Broken model_dump()

Author: rahul-tuli

src/llmcompressor/recipe/recipe.py

@@ -35,6 +35,12 @@ class Recipe(RecipeBase):
     when serializing a recipe, yaml will be used by default.
     """
 
+    version: Optional[str] = Field(default=None)
+    args: RecipeArgs = Field(default_factory=RecipeArgs)
+    stages: List[RecipeStage] = Field(default_factory=list)
+    metadata: Optional[RecipeMetaData] = Field(default=None)
+    args_evaluated: RecipeArgs = Field(default_factory=RecipeArgs)
+
     @classmethod
     def from_modifiers(
         cls,
@@ -280,12 +286,6 @@ def simplify_combine_recipes(
 
         return combined
 
-    version: Optional[str] = None
-    args: RecipeArgs = Field(default_factory=RecipeArgs)
-    stages: List[RecipeStage] = Field(default_factory=list)
-    metadata: Optional[RecipeMetaData] = None
-    args_evaluated: RecipeArgs = Field(default_factory=RecipeArgs)
-
     def calculate_start(self) -> int:
         """
         Calculate and return the start epoch of the recipe.
@@ -507,52 +507,60 @@ def combine_metadata(self, metadata: Optional[RecipeMetaData]):
 
     def model_dump(self, *args, **kwargs) -> Dict[str, Any]:
         """
-        :return: A dictionary representation of the recipe
-        """
-        dict_ = super().model_dump(*args, **kwargs)
-        stages = {}
-
-        for stage in dict_["stages"]:
-            name = f"{stage['group']}_stage"
-            del stage["group"]
+        Generate a serializable dictionary representation of this recipe.
 
-            if name not in stages:
-                stages[name] = []
+        This method transforms the internal recipe structure into a format
+        suitable for YAML serialization while preserving all necessary
+        information for round-trip deserialization.
 
-            stages[name].append(stage)
-
-        dict_["stages"] = stages
+        :param args: Additional positional arguments for parent method
+        :param kwargs: Additional keyword arguments for parent method
+        :return: Dictionary ready for YAML serialization
+        """
+        # Retrieve base representation from parent class
+        raw_dict = super().model_dump(*args, **kwargs)
+
+        # Initialize clean output dictionary
+        serializable_dict = {}
+
+        # Copy recipe metadata attributes
+        metadata_keys = ["version", "args", "metadata"]
+        for key in metadata_keys:
+            if value := raw_dict.get(key):
+                serializable_dict[key] = value
+
+        # Process and organize stages by group
+        if "stages" in raw_dict:
+            # Group stages by their type (e.g., "train", "eval")
+            grouped_stages = {}
+            for stage in raw_dict["stages"]:
+                group_id = (
+                    f"{stage.pop('group')}_stage"  # Remove group field and use as key
+                )
 
-        yaml_recipe_dict = {}
+                if group_id not in grouped_stages:
+                    grouped_stages[group_id] = []
 
-        # populate recipe level attributes
-        recipe_level_attributes = ["version", "args", "metadata"]
+                grouped_stages[group_id].append(stage)
 
-        for attribute in recipe_level_attributes:
-            if attribute_value := dict_.get(attribute):
-                yaml_recipe_dict[attribute] = attribute_value
+            # Format each stage for YAML output
+            for group_id, stages in grouped_stages.items():
+                for idx, stage_data in enumerate(stages):
+                    # Create unique identifiers for multiple stages of same type
+                    final_id = f"{group_id}_{idx}" if len(stages) > 1 else group_id
 
-        # populate stages
-        stages = dict_.pop("stages", {})
-        for stage_name, stage_list in stages.items():
-            for idx, stage in enumerate(stage_list):
-                if len(stage_list) > 1:
-                    # resolve name clashes caused by combining recipes with
-                    # duplicate stage names
-                    final_stage_name = f"{stage_name}_{idx}"
-                else:
-                    final_stage_name = stage_name
-                stage_dict = get_yaml_serializable_stage_dict(
-                    modifiers=stage["modifiers"]
-                )
+                    # Create clean stage representation
+                    stage_yaml = get_yaml_serializable_stage_dict(
+                        modifiers=stage_data["modifiers"]
+                    )
 
-                # infer run_type from stage
-                if run_type := stage.get("run_type"):
-                    stage_dict["run_type"] = run_type
+                    # Preserve run type if specified
+                    if run_type := stage_data.get("run_type"):
+                        stage_yaml["run_type"] = run_type
 
-                yaml_recipe_dict[final_stage_name] = stage_dict
+                    serializable_dict[final_id] = stage_yaml
 
-        return yaml_recipe_dict
+        return serializable_dict
 
     def yaml(self, file_path: Optional[str] = None) -> str:
         """

tests/llmcompressor/recipe/test_recipe.py

@@ -1,107 +1,185 @@
+import os
 import tempfile
 
 import pytest
 import yaml
 
 from llmcompressor.modifiers.obcq.base import SparseGPTModifier
+from llmcompressor.modifiers.pruning.constant import ConstantPruningModifier
 from llmcompressor.recipe import Recipe
 from tests.llmcompressor.helpers import valid_recipe_strings
 
 
 @pytest.mark.parametrize("recipe_str", valid_recipe_strings())
-def test_recipe_create_instance_accepts_valid_recipe_string(recipe_str):
-    recipe = Recipe.create_instance(recipe_str)
-    assert recipe is not None, "Recipe could not be created from string"
+class TestRecipeWithStrings:
+    """Tests that use various recipe strings for validation."""
+
+    def test_create_from_string(self, recipe_str):
+        """Test creating a Recipe from a YAML string."""
+        recipe = Recipe.create_instance(recipe_str)
+        assert recipe is not None, "Recipe could not be created from string"
+        assert isinstance(recipe, Recipe), "Created object is not a Recipe instance"
+
+    def test_create_from_file(self, recipe_str):
+        """Test creating a Recipe from a YAML file."""
+        content = yaml.safe_load(recipe_str)
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") as f:
+            yaml.dump(content, f)
+            f.flush()  # Ensure content is written
+            recipe = Recipe.create_instance(f.name)
+            assert recipe is not None, "Recipe could not be created from file"
+            assert isinstance(recipe, Recipe), "Created object is not a Recipe instance"
+
+    def test_yaml_serialization_roundtrip(self, recipe_str):
+        """
+        Test that a recipe can be serialized to YAML
+        and deserialized back with all properties preserved.
+        """
+        # Create original recipe
+        original_recipe = Recipe.create_instance(recipe_str)
 
+        # Serialize to YAML
+        yaml_str = original_recipe.yaml()
+        assert yaml_str, "Serialized YAML string should not be empty"
 
-@pytest.mark.parametrize("recipe_str", valid_recipe_strings())
-def test_recipe_create_instance_accepts_valid_recipe_file(recipe_str):
-    content = yaml.safe_load(recipe_str)
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml") as f:
-        yaml.dump(content, f)
-        recipe = Recipe.create_instance(f.name)
-        assert recipe is not None, "Recipe could not be created from file"
+        # Deserialize from YAML
+        deserialized_recipe = Recipe.create_instance(yaml_str)
 
+        # Compare serialized forms
+        original_dict = original_recipe.model_dump()
+        deserialized_dict = deserialized_recipe.model_dump()
 
-@pytest.mark.parametrize("recipe_str", valid_recipe_strings())
-def test_yaml_serialization(recipe_str):
-    recipe_instance = Recipe.create_instance(recipe_str)
-    serialized_recipe = recipe_instance.yaml()
-    recipe_from_serialized = Recipe.create_instance(serialized_recipe)
+        assert original_dict == deserialized_dict, "Serialization roundtrip failed"
 
-    expected_dict = recipe_instance.model_dump()
-    actual_dict = recipe_from_serialized.model_dump()
+    def test_model_dump_and_validate(self, recipe_str):
+        """
+        Test that model_dump produces a format compatible
+        with model_validate.
+        """
+        recipe = Recipe.create_instance(recipe_str)
+        validated_recipe = Recipe.model_validate(recipe.model_dump())
+        assert (
+            recipe == validated_recipe
+        ), "Recipe instance and validated recipe do not match"
 
-    assert expected_dict == actual_dict
 
+class TestRecipeSerialization:
+    """
+    Tests for Recipe serialization and deserialization
+      edge cases."""
 
-@pytest.mark.parametrize("recipe_str", valid_recipe_strings())
-def test_model_dump_and_validate(recipe_str):
-    recipe_instance = Recipe.create_instance(recipe_str)
-    validated_recipe = Recipe.model_validate(recipe_instance.model_dump())
-    assert (
-        recipe_instance == validated_recipe
-    ), "Recipe instance and validated recipe do not match"
-
-
-def test_recipe_creates_correct_modifier():
-    start = 1
-    end = 10
-    targets = "__ALL_PRUNABLE__"
-
-    yaml_str = f"""
-        test_stage:
-            pruning_modifiers:
-                ConstantPruningModifier:
-                    start: {start}
-                    end: {end}
-                    targets: {targets}
-        """
+    def test_empty_recipe_serialization(self):
+        """Test serialization of a minimal recipe with no stages."""
+        recipe = Recipe()
+        assert len(recipe.stages) == 0, "New recipe should have no stages"
+
+        # Test roundtrip serialization
+        dumped = recipe.model_dump()
+        loaded = Recipe.model_validate(dumped)
+        assert recipe == loaded, "Empty recipe serialization failed"
 
-    recipe_instance = Recipe.create_instance(yaml_str)
-
-    stage_modifiers = recipe_instance.create_modifier()
-    assert len(stage_modifiers) == 1
-    assert len(modifiers := stage_modifiers[0].modifiers) == 1
-    from llmcompressor.modifiers.pruning.constant import ConstantPruningModifier
-
-    assert isinstance(modifier := modifiers[0], ConstantPruningModifier)
-    assert modifier.start == start
-    assert modifier.end == end
-
-
-def test_recipe_can_be_created_from_modifier_instances():
-    modifier = SparseGPTModifier(
-        sparsity=0.5,
-    )
-    group_name = "dummy"
-
-    # for pep8 compliance
-    recipe_str = (
-        f"{group_name}_stage:\n"
-        "   pruning_modifiers:\n"
-        "       SparseGPTModifier:\n"
-        "           sparsity: 0.5\n"
-    )
-
-    expected_recipe_instance = Recipe.create_instance(recipe_str)
-    expected_modifiers = expected_recipe_instance.create_modifier()
-
-    actual_recipe_instance = Recipe.create_instance(
-        [modifier], modifier_group_name=group_name
-    )
-    actual_modifiers = actual_recipe_instance.create_modifier()
-
-    # assert num stages is the same
-    assert len(actual_modifiers) == len(expected_modifiers)
-
-    # assert num modifiers in each stage is the same
-    assert len(actual_modifiers[0].modifiers) == len(expected_modifiers[0].modifiers)
-
-    # assert modifiers in each stage are the same type
-    # and have the same parameters
-    for actual_modifier, expected_modifier in zip(
-        actual_modifiers[0].modifiers, expected_modifiers[0].modifiers
-    ):
-        assert isinstance(actual_modifier, type(expected_modifier))
-        assert actual_modifier.model_dump() == expected_modifier.model_dump()
+    def test_file_serialization(self):
+        """Test serializing a recipe to a file and reading it back."""
+        recipe = Recipe.create_instance(valid_recipe_strings()[0])
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_path = os.path.join(temp_dir, "recipe.yaml")
+
+            # Write to file
+            recipe.yaml(file_path=file_path)
+            assert os.path.exists(file_path), "YAML file was not created"
+            assert os.path.getsize(file_path) > 0, "YAML file is empty"
+
+            # Read back from file
+            loaded_recipe = Recipe.create_instance(file_path)
+            assert (
+                recipe == loaded_recipe
+            ), "Recipe loaded from file doesn't match original"
+
+
+class TestRecipeModifiers:
+    """Tests for creating and working with modifiers in recipes."""
+
+    def test_creates_correct_modifier(self):
+        """
+        Test that a recipe creates the expected modifier type
+        with correct parameters.
+        """
+        # Recipe parameters
+        params = {"start": 1, "end": 10, "targets": "__ALL_PRUNABLE__"}
+
+        # Create recipe from YAML
+        yaml_str = f"""
+            test_stage:
+                pruning_modifiers:
+                    ConstantPruningModifier:
+                        start: {params['start']}
+                        end: {params['end']}
+                        targets: {params['targets']}
+            """
+        recipe = Recipe.create_instance(yaml_str)
+
+        # Get modifiers from recipe
+        stage_modifiers = recipe.create_modifier()
+        assert len(stage_modifiers) == 1, "Expected exactly one stage modifier"
+
+        modifiers = stage_modifiers[0].modifiers
+        assert len(modifiers) == 1, "Expected exactly one modifier in the stage"
+
+        # Verify modifier type and parameters
+        modifier = modifiers[0]
+        assert isinstance(
+            modifier, ConstantPruningModifier
+        ), "Wrong modifier type created"
+        assert modifier.start == params["start"], "Modifier start value incorrect"
+        assert modifier.end == params["end"], "Modifier end value incorrect"
+        assert modifier.targets == params["targets"], "Modifier targets incorrect"
+
+    def test_create_from_modifier_instances(self):
+        """Test creating a recipe from modifier instances."""
+        # Create a modifier instance
+        sparsity_value = 0.5
+        modifier = SparseGPTModifier(sparsity=sparsity_value)
+        group_name = "dummy"
+
+        # Expected YAML representation
+        recipe_str = (
+            f"{group_name}_stage:\n"
+            "   pruning_modifiers:\n"
+            "       SparseGPTModifier:\n"
+            f"           sparsity: {sparsity_value}\n"
+        )
+
+        # Create recipes for comparison
+        expected_recipe = Recipe.create_instance(recipe_str)
+        actual_recipe = Recipe.create_instance(
+            [modifier], modifier_group_name=group_name
+        )
+
+        # Compare recipes by creating and checking their modifiers
+        self._compare_recipe_modifiers(actual_recipe, expected_recipe)
+
+    def _compare_recipe_modifiers(self, actual_recipe, expected_recipe):
+        """Helper method to compare modifiers created from two recipes."""
+        actual_modifiers = actual_recipe.create_modifier()
+        expected_modifiers = expected_recipe.create_modifier()
+
+        # Compare stage counts
+        assert len(actual_modifiers) == len(expected_modifiers), "Stage counts differ"
+
+        if not actual_modifiers:
+            return  # No modifiers to compare
+
+        # Compare modifier counts in each stage
+        assert len(actual_modifiers[0].modifiers) == len(
+            expected_modifiers[0].modifiers
+        ), "Modifier counts differ"
+
+        # Compare modifier types and parameters
+        for actual_mod, expected_mod in zip(
+            actual_modifiers[0].modifiers, expected_modifiers[0].modifiers
+        ):
+            assert isinstance(actual_mod, type(expected_mod)), "Modifier types differ"
+            assert (
+                actual_mod.model_dump() == expected_mod.model_dump()
+            ), "Modifier parameters differ"