⚡️ Speed up function aggregate_bagging by 40%

[codeflash-ai]

📄 40% (0.40x) speedup for aggregate_bagging in framework/py/flwr/serverapp/strategy/strategy_utils.py

⏱️ Runtime : 13.1 milliseconds → 9.35 milliseconds (best of 275 runs)

📝 Explanation and details

The optimization achieves a 39% speedup by eliminating redundant JSON parsing operations.

Key optimization: The original code called _get_tree_nums() twice - once for each input model - which meant parsing the JSON data four times total (twice in the main function, twice in the helper). The optimized version extracts tree numbers directly from the already-parsed JSON objects, reducing JSON parsing from 4 operations to just 2.

Why this matters: From the line profiler results, JSON parsing (json.loads(bytearray(...))) was the most expensive operation, consuming 25.9% of total runtime in the original version. By eliminating the redundant calls to _get_tree_nums(), we cut this overhead significantly.

Performance characteristics: This optimization is most effective for larger models where JSON parsing overhead dominates. The test cases show benefits across all scales - from single trees to 500+ tree models - because the fixed cost of redundant parsing is eliminated regardless of model size. The optimization maintains identical behavior while reducing the computational complexity from O(4 * JSON_size) to O(2 * JSON_size) for the parsing phase.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import json

# imports
import pytest  # used for our unit tests
from serverapp.strategy.strategy_utils import aggregate_bagging

# function to test
# (see above for aggregate_bagging and _get_tree_nums definition)

# ------------------- Unit Tests for aggregate_bagging -------------------

# Helper functions for creating mock XGBoost JSON models as bytes

def make_xgb_json_bytes(num_trees, num_parallel_tree, start_tree_id=0):
    """
    Create a mock XGBoost model as bytes, with the specified number of trees and parallel trees.
    """
    trees = []
    for i in range(num_trees):
        trees.append({"id": start_tree_id + i, "dummy": True})
    model = {
        "learner": {
            "gradient_booster": {
                "model": {
                    "gbtree_model_param": {
                        "num_trees": str(num_trees),
                        "num_parallel_tree": str(num_parallel_tree),
                    },
                    "iteration_indptr": [0, num_trees],
                    "trees": trees,
                    "tree_info": [0] * num_trees,
                }
            }
        }
    }
    return bytes(json.dumps(model), "utf-8")

def make_xgb_json_bytes_custom(num_trees, num_parallel_tree, iteration_indptr, tree_ids):
    """
    Create a mock XGBoost model as bytes, with custom iteration_indptr and tree IDs.
    """
    trees = []
    for i in tree_ids:
        trees.append({"id": i, "dummy": True})
    model = {
        "learner": {
            "gradient_booster": {
                "model": {
                    "gbtree_model_param": {
                        "num_trees": str(num_trees),
                        "num_parallel_tree": str(num_parallel_tree),
                    },
                    "iteration_indptr": iteration_indptr,
                    "trees": trees,
                    "tree_info": [0] * num_trees,
                }
            }
        }
    }
    return bytes(json.dumps(model), "utf-8")

# ------------------- Basic Test Cases -------------------

def test_aggregate_bagging_empty_prev_returns_curr():
    """If bst_prev_org is empty, return bst_curr_org unchanged."""
    bst_curr = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    codeflash_output = aggregate_bagging(b"", bst_curr); result = codeflash_output

def test_aggregate_bagging_simple_append():
    """Append two trees from bst_curr to two trees from bst_prev, updating tree IDs and counts."""
    bst_prev = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    bst_curr = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    # All tree IDs should be unique and correct
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_curr_parallel_tree_gt_prev():
    """Test where bst_curr has more parallel trees than bst_prev."""
    bst_prev = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    bst_curr = make_xgb_json_bytes(num_trees=3, num_parallel_tree=3)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_curr_parallel_tree_lt_prev():
    """Test where bst_curr has fewer parallel trees than bst_prev."""
    bst_prev = make_xgb_json_bytes(num_trees=3, num_parallel_tree=3)
    bst_curr = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

# ------------------- Edge Test Cases -------------------

def test_aggregate_bagging_prev_has_no_trees():
    """Test when previous model has zero trees."""
    bst_prev = make_xgb_json_bytes(num_trees=0, num_parallel_tree=0)
    bst_curr = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_curr_has_no_trees():
    """Test when current model has zero trees."""
    bst_prev = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    bst_curr = make_xgb_json_bytes(num_trees=0, num_parallel_tree=0)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_prev_and_curr_have_no_trees():
    """Test when both previous and current models have zero trees."""
    bst_prev = make_xgb_json_bytes(num_trees=0, num_parallel_tree=0)
    bst_curr = make_xgb_json_bytes(num_trees=0, num_parallel_tree=0)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]

def test_aggregate_bagging_nonzero_iteration_indptr():
    """Test with custom iteration_indptr (not [0, num_trees])."""
    bst_prev = make_xgb_json_bytes_custom(num_trees=2, num_parallel_tree=2, iteration_indptr=[0, 1, 2], tree_ids=[0, 1])
    bst_curr = make_xgb_json_bytes(num_trees=2, num_parallel_tree=2)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    # tree IDs should be unique and incremented
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_tree_ids_are_sequential():
    """Test that tree IDs are sequential after aggregation, even if curr tree IDs are not."""
    # prev: tree ids [10, 11], curr: tree ids [0, 1] (should be overwritten to [12, 13])
    bst_prev = make_xgb_json_bytes_custom(num_trees=2, num_parallel_tree=2, iteration_indptr=[0, 2], tree_ids=[10, 11])
    bst_curr = make_xgb_json_bytes_custom(num_trees=2, num_parallel_tree=2, iteration_indptr=[0, 2], tree_ids=[0, 1])
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_handles_non_utf8_bytes():
    """Test with bytes containing non-ASCII but valid UTF-8 (should not raise)."""
    # Add a unicode character in a dummy field
    model = {
        "learner": {
            "gradient_booster": {
                "model": {
                    "gbtree_model_param": {
                        "num_trees": "1",
                        "num_parallel_tree": "1",
                    },
                    "iteration_indptr": [0, 1],
                    "trees": [{"id": 0, "dummy": "ü"}],
                    "tree_info": [0],
                }
            }
        }
    }
    bst_prev = bytes(json.dumps(model), "utf-8")
    bst_curr = make_xgb_json_bytes(num_trees=1, num_parallel_tree=1)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    # Should not raise and should aggregate correctly
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

# ------------------- Large Scale Test Cases -------------------

def test_aggregate_bagging_large_scale():
    """Test aggregation with large models (up to 500 trees each)."""
    n = 500
    bst_prev = make_xgb_json_bytes(num_trees=n, num_parallel_tree=n)
    bst_curr = make_xgb_json_bytes(num_trees=n, num_parallel_tree=n)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    # Check all tree IDs are correct
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_large_scale_curr_zero():
    """Test aggregation with large prev model and zero curr model."""
    n = 500
    bst_prev = make_xgb_json_bytes(num_trees=n, num_parallel_tree=n)
    bst_curr = make_xgb_json_bytes(num_trees=0, num_parallel_tree=0)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_large_scale_prev_zero():
    """Test aggregation with zero prev model and large curr model."""
    n = 500
    bst_prev = make_xgb_json_bytes(num_trees=0, num_parallel_tree=0)
    bst_curr = make_xgb_json_bytes(num_trees=n, num_parallel_tree=n)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]

def test_aggregate_bagging_large_scale_nonstandard_ids():
    """Test aggregation with large models and nonstandard starting tree IDs."""
    n = 500
    bst_prev = make_xgb_json_bytes_custom(num_trees=n, num_parallel_tree=n, iteration_indptr=[0, n], tree_ids=list(range(1000, 1000 + n)))
    bst_curr = make_xgb_json_bytes(num_trees=n, num_parallel_tree=n)
    codeflash_output = aggregate_bagging(bst_prev, bst_curr); result = codeflash_output
    result_json = json.loads(result)
    model = result_json["learner"]["gradient_booster"]["model"]
    ids = [tree["id"] for tree in model["trees"]]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import json

# imports
import pytest
from serverapp.strategy.strategy_utils import aggregate_bagging

# ------------------------
# Unit Tests for aggregate_bagging
# ------------------------

# Helper function to generate a minimal valid XGBoost-like model in JSON bytes
def make_xgb_model(num_trees, num_parallel_tree, tree_id_start=0):
    """Creates a minimal XGBoost model dict and returns its bytes encoding."""
    trees = []
    for i in range(num_trees):
        trees.append({"id": tree_id_start + i, "dummy_field": f"tree_{tree_id_start+i}"})
    model = {
        "learner": {
            "gradient_booster": {
                "model": {
                    "gbtree_model_param": {
                        "num_trees": str(num_trees),
                        "num_parallel_tree": str(num_parallel_tree),
                    },
                    "iteration_indptr": [0, num_trees],
                    "trees": trees,
                    "tree_info": [0] * num_trees,
                }
            }
        }
    }
    return bytes(json.dumps(model), "utf-8")

# Helper to extract model details from bytes
def extract_model_details(model_bytes):
    model = json.loads(model_bytes.decode("utf-8"))
    model_data = model["learner"]["gradient_booster"]["model"]
    num_trees = int(model_data["gbtree_model_param"]["num_trees"])
    tree_ids = [tree["id"] for tree in model_data["trees"]]
    iteration_indptr = model_data["iteration_indptr"]
    tree_info = model_data["tree_info"]
    return num_trees, tree_ids, iteration_indptr, tree_info

# ------------------------
# 1. Basic Test Cases
# ------------------------

def test_aggregate_bagging_prev_empty_returns_current():
    """If previous model is empty, current model should be returned unchanged."""
    curr = make_xgb_model(2, 2)
    codeflash_output = aggregate_bagging(b"", curr); result = codeflash_output

def test_aggregate_bagging_single_tree_each():
    """Aggregate two models, each with one tree."""
    prev = make_xgb_model(1, 1)
    curr = make_xgb_model(1, 1, tree_id_start=0)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_multiple_trees():
    """Aggregate prev with 2 trees and curr with 3 trees."""
    prev = make_xgb_model(2, 2)
    curr = make_xgb_model(3, 3)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_tree_id_offset():
    """Ensure tree IDs in curr are offset by prev's num_trees."""
    prev = make_xgb_model(2, 2)
    curr = make_xgb_model(2, 2, tree_id_start=100)  # IDs in curr ignored
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    _, tree_ids, _, _ = extract_model_details(result)

# ------------------------
# 2. Edge Test Cases
# ------------------------

def test_aggregate_bagging_curr_zero_trees():
    """If curr has zero trees, prev should remain unchanged except for iteration_indptr."""
    prev = make_xgb_model(2, 2)
    # curr with zero trees
    curr = make_xgb_model(0, 0)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_prev_zero_trees():
    """Prev has zero trees, curr has some trees."""
    prev = make_xgb_model(0, 0)
    curr = make_xgb_model(3, 3)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_curr_and_prev_zero_trees():
    """Both prev and curr have zero trees."""
    prev = make_xgb_model(0, 0)
    curr = make_xgb_model(0, 0)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_curr_tree_ids_are_ignored():
    """Tree IDs in curr are ignored and replaced by correct offset."""
    prev = make_xgb_model(3, 3)
    # curr trees have negative IDs, which should be overwritten
    curr_model = {
        "learner": {
            "gradient_booster": {
                "model": {
                    "gbtree_model_param": {
                        "num_trees": "2",
                        "num_parallel_tree": "2",
                    },
                    "iteration_indptr": [0, 2],
                    "trees": [{"id": -1}, {"id": -2}],
                    "tree_info": [0, 0],
                }
            }
        }
    }
    curr_bytes = bytes(json.dumps(curr_model), "utf-8")
    codeflash_output = aggregate_bagging(prev, curr_bytes); result = codeflash_output
    _, tree_ids, _, _ = extract_model_details(result)

def test_aggregate_bagging_iteration_indptr_grows_correctly():
    """iteration_indptr should grow by the number of trees in curr."""
    prev = make_xgb_model(4, 4)
    curr = make_xgb_model(2, 2)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    _, _, iteration_indptr, _ = extract_model_details(result)

def test_aggregate_bagging_tree_info_grows_correctly():
    """tree_info should be extended with zeros for each new tree."""
    prev = make_xgb_model(2, 2)
    curr = make_xgb_model(3, 3)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    _, _, _, tree_info = extract_model_details(result)

def test_aggregate_bagging_non_utf8_bytes():
    """Should raise a JSONDecodeError if bytes are not valid JSON."""
    prev = b"\x80\x81\x82"
    curr = make_xgb_model(1, 1)
    with pytest.raises(Exception):
        aggregate_bagging(prev, curr)

def test_aggregate_bagging_missing_fields():
    """Should raise KeyError if required fields are missing."""
    prev = make_xgb_model(1, 1)
    # curr missing 'trees'
    curr_model = {
        "learner": {
            "gradient_booster": {
                "model": {
                    "gbtree_model_param": {
                        "num_trees": "1",
                        "num_parallel_tree": "1",
                    },
                    "iteration_indptr": [0, 1],
                    # "trees": missing
                    "tree_info": [0],
                }
            }
        }
    }
    curr_bytes = bytes(json.dumps(curr_model), "utf-8")
    with pytest.raises(KeyError):
        aggregate_bagging(prev, curr_bytes)

# ------------------------
# 3. Large Scale Test Cases
# ------------------------

def test_aggregate_bagging_large_scale():
    """Aggregate two large models and check correctness and performance."""
    prev = make_xgb_model(500, 500)
    curr = make_xgb_model(400, 400)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_large_curr_small_prev():
    """Aggregate small prev with large curr."""
    prev = make_xgb_model(1, 1)
    curr = make_xgb_model(999, 999)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_large_prev_small_curr():
    """Aggregate large prev with small curr."""
    prev = make_xgb_model(999, 999)
    curr = make_xgb_model(1, 1)
    codeflash_output = aggregate_bagging(prev, curr); result = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(result)

def test_aggregate_bagging_many_small_merges():
    """Repeatedly aggregate small models to test cumulative effects."""
    prev = make_xgb_model(0, 0)
    for i in range(10):
        curr = make_xgb_model(100, 100)
        codeflash_output = aggregate_bagging(prev, curr); prev = codeflash_output
    num_trees, tree_ids, iteration_indptr, tree_info = extract_model_details(prev)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from serverapp.strategy.strategy_utils import aggregate_bagging
import pytest

def test_aggregate_bagging():
    with pytest.raises(JSONDecodeError, match='Expecting\\ value:\\ line\\ 1\\ column\\ 2\\ \\(char\\ 1\\)'):
        aggregate_bagging(b' ', b'')

def test_aggregate_bagging_2():
    aggregate_bagging(v1:=b'', v1)

🔎 Concolic Coverage Tests and Runtime

To edit these changes git checkout codeflash/optimize-aggregate_bagging-mh9kfwe2 and push.