⚡️ Speed up method CompletionTrainingParametersIn.serialize_model by 17%

[codeflash-ai]

📄 17% (0.17x) speedup for CompletionTrainingParametersIn.serialize_model in src/mistralai/models/completiontrainingparametersin.py

⏱️ Runtime : 1.01 milliseconds → 857 microseconds (best of 34 runs)

📝 Explanation and details

The optimized code achieves a 17% speedup through several key micro-optimizations that reduce redundant operations in the serialization loop:

Key optimizations:

1. Hoisted expensive lookups: Moved self.__pydantic_fields_set__ and type(self).model_fields outside the loop to avoid repeated attribute access on each iteration.

2. Optimized dictionary operations: Replaced the pattern of serialized.get(k) followed by serialized.pop(k, None) with a single serialized.pop(k, None) call, eliminating one dictionary lookup per field.

3. Reduced attribute access: Created local variables (pydantic_fields_set, model_fields) to minimize dot notation lookups within the hot loop.

4. Style improvement: Changed not k in optional_fields to k not in optional_fields for better readability (though minimal performance impact).

Why it's faster:

• Dictionary operations (get/pop) and attribute access (self.__pydantic_fields_set__) are relatively expensive in Python's interpreter
• The loop iterates over 7 fields, so each micro-optimization compounds
• Hoisting invariant computations outside loops is a classic performance optimization

Test case performance:
The optimizations show consistent 7-18% improvements across all test scenarios, with particularly strong gains in:

• Large-scale tests with many instances (18.1% faster)
• Tests with all fields set to values (13.5% faster)
• Default/unset field scenarios (10.7% faster)

The optimizations are most effective for workloads that serialize many instances, as the per-instance overhead reduction compounds significantly.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 1616 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

from typing import Optional

# imports
import pytest
from mistralai.models.completiontrainingparametersin import \
    CompletionTrainingParametersIn


# Simulate the types and constants used in the source code
class UNSET_SENTINEL_TYPE:
    pass
UNSET_SENTINEL = UNSET_SENTINEL_TYPE()
UNSET = UNSET_SENTINEL

# Simulate BaseModel and model_serializer decorator
class BaseModel:
    # Simulate Pydantic's __pydantic_fields_set__ to track explicitly set fields
    def __init__(self, **kwargs):
        self.__pydantic_fields_set__ = set(kwargs.keys())
        for k, v in kwargs.items():
            setattr(self, k, v)

    # Simulate model_fields as a dict of field name to field info
    @classmethod
    def model_fields(cls):
        # Each field is a dummy object with .alias attribute
        fields = {}
        for name in getattr(cls, "_fields", []):
            fields[name] = type("Field", (), {"alias": None})()
        return fields

def model_serializer(mode=None):
    # Just a no-op decorator for testing
    def decorator(fn):
        return fn
    return decorator
from mistralai.models.completiontrainingparametersin import \
    CompletionTrainingParametersIn


# Helper function to simulate handler
def default_handler(obj):
    # Return a dict of field names and their values
    result = {}
    for name in obj._fields:
        result[name] = getattr(obj, name)
    return result

# ------------------- UNIT TESTS -------------------

# Basic Test Cases


















#------------------------------------------------
from __future__ import annotations

from typing import Optional

# imports
import pytest
from mistralai.models.completiontrainingparametersin import \
    CompletionTrainingParametersIn


# Simulate UNSET and UNSET_SENTINEL as in the original code
class _UnsetType:
    pass
UNSET = _UnsetType()
UNSET_SENTINEL = _UnsetType()

# Simulate OptionalNullable for test purposes
def OptionalNullable(type_):
    return Optional[type_]

# Simulate BaseModel for test purposes
class BaseModel:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        # Track which fields are explicitly set
        self.__pydantic_fields_set__ = set(kwargs.keys())
    @classmethod
    def model_fields(cls):
        # Simulate pydantic's model_fields
        return {k: type('', (), {'alias': None})() for k in cls.__annotations__}

# Simulate model_serializer decorator
def model_serializer(mode=None):
    def decorator(f):
        return f
    return decorator
from mistralai.models.completiontrainingparametersin import \
    CompletionTrainingParametersIn


# Helper handler function for serialization
def default_handler(obj):
    # Simulate pydantic's dict() output
    result = {}
    for field in CompletionTrainingParametersIn.model_fields.keys():
        result[field] = getattr(obj, field, UNSET_SENTINEL)
    return result

# -------------------- UNIT TESTS --------------------

# Basic Test Cases


def test_serialize_model_some_fields_unset():
    # Some fields set, others left as UNSET
    params = CompletionTrainingParametersIn(
        training_steps=200,
        learning_rate=0.02,
        # weight_decay not set
        warmup_fraction=0.3,
        # epochs not set
        seq_len=1024,
        # fim_ratio not set
    )
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 8.70μs -> 7.91μs (9.96% faster)

def test_serialize_model_defaults_only():
    # No fields set (all default)
    params = CompletionTrainingParametersIn()
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 8.28μs -> 7.48μs (10.7% faster)

def test_serialize_model_explicit_none_nullable_field():
    # Set a nullable field explicitly to None
    params = CompletionTrainingParametersIn(
        training_steps=None,
        learning_rate=0.1,
    )
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 8.39μs -> 7.33μs (14.5% faster)

def test_serialize_model_explicit_none_non_nullable_field():
    # Set a non-nullable field to None (learning_rate)
    params = CompletionTrainingParametersIn(
        learning_rate=None,
    )
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 8.07μs -> 7.50μs (7.67% faster)

# Edge Test Cases

def test_serialize_model_zero_and_negative_values():
    # Set fields to zero and negative numbers
    params = CompletionTrainingParametersIn(
        training_steps=0,
        learning_rate=-0.01,
        weight_decay=0,
        warmup_fraction=-0.5,
        epochs=0,
        seq_len=-256,
        fim_ratio=0,
    )
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 6.47μs -> 5.90μs (9.67% faster)



def test_serialize_model_only_nullable_fields_set_none():
    # All nullable fields explicitly set to None
    params = CompletionTrainingParametersIn(
        training_steps=None,
        weight_decay=None,
        warmup_fraction=None,
        epochs=None,
        seq_len=None,
        fim_ratio=None,
    )
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 8.33μs -> 7.73μs (7.76% faster)


def test_serialize_model_large_scale_all_fields_set():
    # Simulate large scale by setting all fields to large values
    params = CompletionTrainingParametersIn(
        training_steps=999999,
        learning_rate=999.999,
        weight_decay=999.999,
        warmup_fraction=0.999,
        epochs=999,
        seq_len=999,
        fim_ratio=0.999,
    )
    codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 8.37μs -> 7.37μs (13.5% faster)

def test_serialize_model_large_scale_many_instances():
    # Create and serialize many instances with varying values
    results = []
    for i in range(100):  # 100 instances, each with different values
        params = CompletionTrainingParametersIn(
            training_steps=i,
            learning_rate=i * 0.01,
            weight_decay=i * 0.02,
            warmup_fraction=i * 0.001,
            epochs=i,
            seq_len=i * 10,
            fim_ratio=i * 0.005,
        )
        codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 273μs -> 231μs (18.1% faster)
        results.append(result)

def test_serialize_model_large_scale_nullable_none():
    # Create many instances with all nullable fields set to None
    for i in range(100):
        params = CompletionTrainingParametersIn(
            training_steps=None,
            weight_decay=None,
            warmup_fraction=None,
            epochs=None,
            seq_len=None,
            fim_ratio=None,
            learning_rate=0.0001 + i,
        )
        codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 277μs -> 236μs (17.5% faster)

def test_serialize_model_large_scale_unset_fields():
    # Create many instances with all fields UNSET except learning_rate
    for i in range(100):
        params = CompletionTrainingParametersIn(
            learning_rate=0.0001 + i,
        )
        codeflash_output = params.serialize_model(default_handler); result = codeflash_output # 388μs -> 329μs (18.0% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-CompletionTrainingParametersIn.serialize_model-mh4hujlb and push.