26. Load-Balance Requests Across Model Replicas

Problem

Inference requests of varying cost (short vs. long utterances) arrive one at a time and must be assigned to one of a fixed number of model replicas. Always send each request to whichever replica currently has the least total assigned load, so no single replica gets overloaded while others sit idle.

Report which replica each request lands on, in arrival order, plus each replica's final total load.

Source: src/26_load_balance_replicas.py

def assign_requests_to_replicas(
    num_replicas: int, request_costs: list[int]
) -> tuple[list[int], list[int]]: ...

>>> assign_requests_to_replicas(2, [1, 1, 1, 1])
([0, 1, 0, 1], [2, 2])
>>> assign_requests_to_replicas(1, [3, 4, 5])
([0, 0, 0], [12])

Step-by-Step Approach

1. Seed a min-heap with one entry per replica: (current_load=0, replica_index), so the least-loaded replica is always at the top.
2. For each incoming request cost, pop the top of the heap — that's the currently least-loaded replica — and record that request's assignment as that replica's index.
3. Push the same replica back onto the heap with its load incremented by the request's cost, so the heap always reflects current state for the next request.
4. Rely on tuple comparison to break load ties by the lower replica index automatically, since (load, replica_index) compares the index only when loads are equal.
5. After processing every request, drain the heap into a final_loads array indexed by replica number for reporting.

The key insight is that "always route to the currently least-loaded option" is a greedy that's optimal for minimizing the maximum load one request at a time, and a min-heap keyed on current load is exactly the structure that answers "which is least-loaded right now" in O(log r) instead of scanning all r replicas on every request.

Reference solution

import heapq


def assign_requests_to_replicas(
    num_replicas: int, request_costs: list[int]
) -> tuple[list[int], list[int]]:
    # Min-heap of (current_load, replica_index): always pop the least-loaded
    # replica, assign the request, add its cost, and push it back. Tuple
    # comparison naturally tie-breaks on the lower index when loads match.
    # O(n log r) time for n requests across r replicas.
    heap = [(0, r) for r in range(num_replicas)]
    heapq.heapify(heap)

    assignment: list[int] = []
    for cost in request_costs:
        load, replica = heapq.heappop(heap)
        assignment.append(replica)
        heapq.heappush(heap, (load + cost, replica))

    final_loads = [0] * num_replicas
    for load, replica in heap:
        final_loads[replica] = load

    return assignment, final_loads

Key Functions & Tricks

• heapq.heapify — builds the initial all-zero-load heap in O(r).
• (load, replica) tuple ordering — sorts by load first, breaking ties on the lower replica index for free.
• Pop-modify-push per request — the standard heap idiom for "always act on the current minimum, then update it," O(log r) per request.
• Draining the heap into an index-ordered final_loads array at the end — heap order isn't replica order, so this reconstructs a stable, readable report.

How to Recognize This Pattern