Apache Kafka Partition Assignments¶
Problem Statement¶
Optimize the performance of the Apache Kafka cluster by smartly assigning Topic partitions to the various brokers. The goal is to minimize the latency between the different brokers (for, for instance, replication or election workloads) while dividing the partitions somewhat evenly over the different brokers.
Full example scenario
As a Data Engineer at a software company, you are tasked with optimizing the performance of your Apache Kafka cluster. Kafka topics are used to manage data streams, and each topic is divided into partitions. Efficient partitioning of these topics can significantly reduce latency and improve throughput. Given data about the cluster’s brokers and topics, which number of partitions should each topic have, and which brokers should host them, in order to maximize throughput by minimizing data transfer between brokers? For every topic, this transfer cost is calculated as the average transfer rate between all broker pairs, weighted by how many partitions of that topic each one holds.
The number of partitions assigned to each broker must not exceed that broker’s maximum capacity.
The number of partitions assigned to each topic must lie within its specified minimum and maximum range, and every topic must be assigned.
The total number of partitions assigned to each broker should be as balanced as possible, with no broker deviating more than 20% from the broker average.
For each topic, no broker can hold more than 40% of that topic’s partitions.
The brokers can be found in
kafka_brokers.csv, with the columnsBroker_ID,Max_Partitions, the topics can be found inkafka_topics.csv, with the columnsTopic_ID,Min_Partitions,Max_Partitions, and the data transfer rates can be found inkafka_rates.csv, with the columnsFrom_Broker,To_Broker,Transfer_Rate.
What an optimal solution looks like
Built with Modeler¶
The problem above is stated as a single shot: you hand over the whole problem and get a model back. Modeler takes the other route — you start with the business question you actually have, and it works with you until the problem is pinned down, the data is in, the code runs, and the answer holds up when the problem changes.
The tabs below walk the same Kafka partition-assignment problem through that route, one Modeler capability at a time.
Instead of writing the whole prompt up front, you start with the business question you actually have. Modeler searches the knowledge base, asks one clarifying question at a time — always the most important one first — and tells you which parts of the model are still missing. The options it offers come from examples in the knowledge base and from how similar problems are usually modeled in your industry, so you can answer with a number instead of prose.
I’m a data engineer at a software company running an Apache Kafka cluster. Each topic is split into partitions, and I need to decide how many partitions each topic should have and which brokers to place them on, so that inter-broker network traffic (from things like replication) stays low and no single broker gets overloaded. Can you help me build this as an optimization model?
This reads as an assignment problem — partitions of each topic are the items, brokers are the bins — with a traffic-minimization objective instead of a simple count. Let me check the knowledge base for how similar cluster-placement problems are usually modeled before asking anything.
The core driver of inter-broker traffic in Kafka is replication — each partition has a leader on one broker and replicas on others, and data flows from the leader to each replica. Before anything else, I need to know how precisely the user wants that traffic modeled, since it decides whether replication needs its own variables or can be approximated directly from where partitions land — so that has to be the first question.
Good, I have useful context from the knowledge base. Let me now start building the specification with you by asking questions one at a time.
Question 1 of ~4: Does every partition share the same replication factor, and should the model represent replication explicitly, or approximate inter-broker traffic some other way?
All topics use the same global replication factor, modeled explicitly (e.g., 3 replicas per partition, each a separate placement decision).
Each topic has its own fixed replication factor, modeled explicitly as an input.
Replication is not modeled explicitly — inter-broker traffic is approximated purely from the pairwise transfer rates between brokers that co-host partitions of the same topic.
Which option best matches your situation, or would you like to describe it differently?
Option 3 — no explicit replication factor, just approximate inter-broker traffic from the pairwise transfer rate between two brokers, weighted by how many partitions of that topic each one holds — specifically as the product of the two partition counts, even though that makes the objective quadratic rather than linear; it’s the most faithful description of the actual cost. I have 4 brokers, each with a maximum partition capacity, and 4 topics, each with a minimum and maximum number of partitions it must be split into — every topic must be assigned. Two more things: I want the load balanced across brokers (no broker should end up more than 20% away from the average partition count across brokers), and for any single topic, no one broker should hold more than 40% of that topic’s partitions.
That single answer resolves the traffic-modeling question and volunteers everything else at once: the exact objective weighting (product of partition counts, accepted as quadratic), the topic partition-count ranges, the broker capacity constraint, a 20% load-balance rule, and a 40% per-topic concentration rule. Nothing about the objective or the constraints is still ambiguous — this is a non-convex MIQP, and the specification is buildable as-is.
Got it. Here is what I have captured so far:
That’s a complete specification to build from — no open questions remain. I’ll save this and move on to bringing in your data.
Modeler updates the specification after every answer, even while it is still incomplete — so you can always see what it has understood and what it is still waiting on.
With the specification pinned down, Modeler asked for the broker, topic, and transfer-rate figures as files rather than as numbers typed into the chat. Attach all three CSVs to the conversation and it reads the real shape of your data instead of guessing at it — entity names, column headers, units and all.
Broker_ID,Max_Partitions
B1,50
B2,60
B3,55
B4,65
Topic_ID,Min_Partitions,Max_Partitions
T1,10,20
T2,15,25
T3,8,18
T4,12,22
From_Broker,To_Broker,Transfer_Rate
B1,B2,5
B1,B3,3
B1,B4,4
B2,B1,5
B2,B3,4
B2,B4,7
B3,B1,3
B3,B2,4
B3,B4,2
B4,B1,4
B4,B2,7
B4,B3,2
kafka_brokers.csv, kafka_topics.csv and kafka_rates.csv attached
Here are the three files with real values — brokers, topics, and the pairwise transfer rate matrix (units are relative network cost per partition-pair, higher means costlier to spread that topic across those two brokers).
Reading all three attachments. Brokers and topics are small, one row each; the rates file is a full 4x4 matrix minus the diagonal, already symmetric, so no unit conversion or reconciliation is needed. Worth checking that every broker and topic id referenced in the rates file actually exists in the other two before building from this blind.
The specification is now complete with all your data. Before moving to implementation, here is a quick summary of what has been captured — please confirm or flag anything to adjust:
Are you happy with the specification as-is, or would you like to adjust anything before I generate the implementation?
Where your files live
Everything in the conversation — the files you attached and every artifact Modeler generates from them — lands in the conversation workspace as an attachment. Each one opens in a side panel next to the chat, so you can read the specification while you keep answering questions, or scroll the code while you ask for a change. From that panel you can copy the contents or download the file.
By the end of this conversation the workspace holds six items:
Attachment |
What it is |
|---|---|
|
The broker, topic, and transfer-rate data you attached. |
|
The specification, rewritten after every answer. |
|
The |
|
The |
|
Output of the test runs, including the solver logs. |
|
The mathematical formulation, generated on request. |
Starting from files you already have
The same upload path runs in reverse. Attach the files of an existing
optimization project — a specification, a model script, a test file, or any
subset — and ask Modeler to check the project. It completes the
specification where it is thin (asking clarifying questions, as above), fills
in a missing implementation or test suite, reconciles code that has drifted
from the specification, and runs the tests. Point it at a model written for
another solver or another language and it recovers a specification from that
file, asks you to confirm it, then reimplements the model in gurobipy.
The specification is the artifact the questions are building toward, and the one everything else is generated from. Modeler rewrites it after every answer — even while it is still incomplete — so it doubles as the running record of what has been agreed and what is still open. It follows a fixed seven-section structure and is written entirely in business terms: no mathematics, so the person who owns the problem can read it and say whether it is right.
1. Overview
A software company runs an Apache Kafka cluster with multiple topics and brokers. Each topic must be split into a number of partitions, and each partition must be assigned to a broker. Inter-broker network traffic — approximated by pairwise transfer rates between brokers weighted by the number of partitions they co-host for the same topic — must be minimized. At the same time, no broker should be overloaded, and no single topic should be concentrated on too few brokers.
2. Objective
Minimize total inter-broker traffic cost, defined as the sum over all topics and all broker pairs of the pairwise transfer rate multiplied by the product of the number of partitions of that topic on each of the two brokers.
3. Constraints
C1 — Topic partition bounds. For each topic, the total number of partitions assigned across all brokers must be at least the topic’s minimum and at most its maximum allowed partition count. Every topic must be assigned.
C2 — Broker capacity. For each broker, the total number of partitions assigned to it, across all topics, must not exceed its maximum partition capacity.
C3 — Global load balance (20% rule). No broker’s total partition count may deviate by more than 20% from the average partition count across all brokers, so no single broker is left carrying a disproportionate share of the cluster.
C4 — Per-topic concentration limit (40% rule). For any single topic, no broker may hold more than 40% of that topic’s total assigned partitions — this prevents any one broker from becoming a hotspot for a given topic.
C5 — Non-negativity and integrality. The number of partitions of each topic assigned to each broker must be a non-negative integer.
4. Input
Entity |
Fields |
Values |
|---|---|---|
Broker |
broker_id, max_capacity |
4 brokers (B1–B4); capacities 50, 60, 55, 65 — 230 partitions of total capacity |
Topic |
topic_id, min_partitions, max_partitions |
4 topics (T1–T4); ranges [10,20], [15,25], [8,18], [12,22] — 45 partitions at the minimum, 85 at the maximum |
Transfer rate |
from_broker, to_broker, transfer_rate |
12 ordered pairs (symmetric); cheapest pair B3–B4 (rate 2), costliest pair B2–B4 (rate 7) |
Total broker capacity (230) comfortably covers even the maximum combined topic demand (85), so capacity is rarely the binding constraint — the real tension is between the load-balance and concentration rules, and the pull toward the cheapest broker pairs.
5. Output
The number of partitions of each topic assigned to each broker.
The total inter-broker traffic cost.
Each topic’s total assigned partitions and each broker’s total load.
6. Test scenarios
ID |
Scenario |
Acceptance criteria |
|---|---|---|
T1 |
Standard optimization — the full 4-broker / 4-topic dataset. |
An optimal solution: every topic’s partitions within its bounds, no broker over capacity, load balance and concentration rules satisfied, and a positive finite objective. |
T2 |
Tight capacity — broker capacities set just high enough to cover the minimum partition requirements (sum of minimums = 45). |
Still optimal and feasible under every constraint, with no slack to spare. |
T3 |
Infeasible — broker capacities (total 20) fall below the sum of the topics’ minimum partition requirements (45). |
Reported as infeasible; no assignment is returned. |
T4 |
Infeasible — only 2 brokers, each capped at 40% of a topic’s partitions, so together they can place at most 80%. |
Reported as infeasible by construction, independent of capacity. |
T5 |
High transfer rate avoidance — one broker pair set to a much higher rate than the rest. |
The objective is strictly higher than an equivalent uniform-rate scenario, proving the rate matrix is genuinely embedded in the cost. |
7. Out of scope
Explicit replication factor modeling (leader/follower assignment is not modeled).
Partition reassignment cost (migration cost from a previous assignment is not considered).
Rack-awareness constraints.
Dynamic or time-varying transfer rates.
With the specification settled, Modeler writes the gurobipy
implementation. This is the version that came out of the run-and-fix loop in
the Cloud solve tab — the model itself was correct from the first run;
what the loop actually fixed was a test expectation, as that tab explains.
"""
Kafka Partition Assignment Optimizer
-------------------------------------
Minimizes inter-broker network traffic in an Apache Kafka cluster by deciding
how many partitions each topic should place on each broker.
The traffic cost is approximated as:
sum over all topics t and all ordered broker pairs (i, j) of
r[i,j] * x[t,i] * x[t,j]
where x[t,i] is the integer number of partitions of topic t on broker i and
r[i,j] is the pairwise transfer rate between brokers i and j.
Because the objective contains products of two integer decision variables,
the model is a non-convex Mixed-Integer Quadratic Program (MIQP).
Gurobi solves it with NonConvex = 2.
"""
import gurobipy as gp
from gurobipy import GRB
import pandas as pd
# ---------------------------------------------------------------------------
# Optimizer class
# ---------------------------------------------------------------------------
class KafkaPartitionOptimizer:
"""
Business optimizer for Kafka partition assignment.
Usage
-----
with gp.Env(empty=True) as env:
env.setParam('OutputFlag', 0)
env.start()
with KafkaPartitionOptimizer(env) as opt:
opt.set_data(brokers_df, topics_df, rates_df)
opt.build_model()
opt.solve()
solution = opt.get_solution()
opt.display_results()
"""
# Load-balance tolerance: each broker's load must be within this fraction
# of the fleet average (20 % rule).
LOAD_BALANCE_TOLERANCE = 0.20
# Maximum share of a topic's partitions any single broker may hold (40 % rule).
CONCENTRATION_LIMIT = 0.40
def __init__(self, env):
"""Initialise the optimizer with a Gurobi environment."""
self.env = env
self.model = None
# Data holders
self.brokers = [] # list of broker_id strings
self.topics = [] # list of topic_id strings
self.max_capacity = {} # broker_id -> int
self.min_partitions = {} # topic_id -> int
self.max_partitions = {} # topic_id -> int
self.transfer_rate = {} # (from_broker, to_broker) -> float
# Decision variables
self._x = {} # (topic_id, broker_id) -> Var
# ------------------------------------------------------------------
# Context manager
# ------------------------------------------------------------------
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if self.model is not None:
self.model.dispose()
return False
# ------------------------------------------------------------------
# Data loading
# ------------------------------------------------------------------
def set_data(self, brokers_df: pd.DataFrame,
topics_df: pd.DataFrame,
rates_df: pd.DataFrame):
"""
Load input data from three DataFrames.
Parameters
----------
brokers_df : columns [broker_id, max_capacity]
topics_df : columns [topic_id, min_partitions, max_partitions]
rates_df : columns [from_broker, to_broker, transfer_rate]
"""
self.brokers = brokers_df["broker_id"].tolist()
self.max_capacity = dict(
zip(brokers_df["broker_id"], brokers_df["max_capacity"])
)
self.topics = topics_df["topic_id"].tolist()
self.min_partitions = dict(
zip(topics_df["topic_id"], topics_df["min_partitions"])
)
self.max_partitions = dict(
zip(topics_df["topic_id"], topics_df["max_partitions"])
)
self.transfer_rate = {
(row["from_broker"], row["to_broker"]): row["transfer_rate"]
for _, row in rates_df.iterrows()
}
# ------------------------------------------------------------------
# Model construction
# ------------------------------------------------------------------
def build_model(self):
"""Build the MIQP model: variables, objective, and all constraints."""
self.model = gp.Model("kafka_partition_assignment", env=self.env)
# Non-convex MIQP requires NonConvex = 2
self.model.setParam("NonConvex", 2)
self._create_variables()
self._set_objective()
self._add_constraints()
def _create_variables(self):
"""
x[t, i] : integer number of partitions of topic t assigned to broker i.
Lower bound is 0 (Gurobi default).
"""
self._x = self.model.addVars(
self.topics,
self.brokers,
vtype=GRB.INTEGER,
lb=0,
name="x"
)
def _set_objective(self):
"""
Minimise total inter-broker traffic cost:
sum_{t} sum_{(i,j): i != j} r[i,j] * x[t,i] * x[t,j]
The transfer_rate dict already contains both directions (symmetric),
so iterating over all (i, j) with i != j captures every ordered pair.
"""
traffic_cost = gp.quicksum(
self.transfer_rate[(i, j)] * self._x[t, i] * self._x[t, j]
for t in self.topics
for i in self.brokers
for j in self.brokers
if i != j and (i, j) in self.transfer_rate
)
self.model.setObjective(traffic_cost, GRB.MINIMIZE)
def _add_constraints(self):
"""Add all business constraints to the model."""
self._add_topic_partition_bounds()
self._add_broker_capacity()
self._add_load_balance()
self._add_concentration_limit()
def _add_topic_partition_bounds(self):
"""
C1 — Topic Partition Bounds:
The total partitions assigned for each topic must lie within
[min_partitions, max_partitions].
"""
for t in self.topics:
topic_total = gp.quicksum(self._x[t, i] for i in self.brokers)
self.model.addConstr(
topic_total >= self.min_partitions[t],
name=f"c1_min_{t}"
)
self.model.addConstr(
topic_total <= self.max_partitions[t],
name=f"c1_max_{t}"
)
def _add_broker_capacity(self):
"""
C2 — Broker Capacity:
The total partitions on each broker (across all topics) must not exceed
its maximum capacity.
"""
for i in self.brokers:
self.model.addConstr(
gp.quicksum(self._x[t, i] for t in self.topics)
<= self.max_capacity[i],
name=f"c2_capacity_{i}"
)
def _add_load_balance(self):
"""
C3 — Global Load Balance (20 % Rule):
Every broker's total partition count must lie within 20 % of the
fleet-wide average. Letting B = number of brokers and
total = sum of all x[t,i], the constraint is:
B * load_i >= 0.8 * total (lower bound)
B * load_i <= 1.2 * total (upper bound)
Both are linear constraints on the integer variables.
"""
B = len(self.brokers)
lb_factor = 1.0 - self.LOAD_BALANCE_TOLERANCE # 0.8
ub_factor = 1.0 + self.LOAD_BALANCE_TOLERANCE # 1.2
total_partitions = gp.quicksum(
self._x[t, i] for t in self.topics for i in self.brokers
)
for i in self.brokers:
broker_load = gp.quicksum(self._x[t, i] for t in self.topics)
self.model.addConstr(
B * broker_load >= lb_factor * total_partitions,
name=f"c3_lb_{i}"
)
self.model.addConstr(
B * broker_load <= ub_factor * total_partitions,
name=f"c3_ub_{i}"
)
def _add_concentration_limit(self):
"""
C4 — Per-Topic Concentration Limit (40 % Rule):
No single broker may hold more than 40 % of any topic's total partitions.
x[t, i] <= 0.4 * sum_{i'} x[t, i'] for all t, i
"""
for t in self.topics:
topic_total = gp.quicksum(self._x[t, i] for i in self.brokers)
for i in self.brokers:
self.model.addConstr(
self._x[t, i] <= self.CONCENTRATION_LIMIT * topic_total,
name=f"c4_conc_{t}_{i}"
)
# ------------------------------------------------------------------
# Solving
# ------------------------------------------------------------------
def solve(self):
"""Optimise the model and return the Gurobi status code."""
self.model.optimize()
return self.model.Status
# ------------------------------------------------------------------
# Solution retrieval
# ------------------------------------------------------------------
def get_solution(self) -> dict:
"""
Extract and return the business solution as a dictionary.
Returns
-------
dict with keys:
status : Gurobi status code (int)
is_optimal : True if an optimal solution was found
objective_value : total traffic cost (float) or None
partition_assignment : {(topic_id, broker_id): count} non-zero only
topic_totals : {topic_id: total partitions assigned}
broker_loads : {broker_id: total partitions hosted}
"""
status = self.model.Status
is_optimal = (status == GRB.OPTIMAL)
if not is_optimal:
return {
"status": status,
"is_optimal": False,
"objective_value": None,
"partition_assignment": {},
"topic_totals": {},
"broker_loads": {},
}
partition_assignment = {
(t, i): int(round(self._x[t, i].X))
for t in self.topics
for i in self.brokers
if self._x[t, i].X > 0.5
}
topic_totals = {
t: sum(int(round(self._x[t, i].X)) for i in self.brokers)
for t in self.topics
}
broker_loads = {
i: sum(int(round(self._x[t, i].X)) for t in self.topics)
for i in self.brokers
}
return {
"status": status,
"is_optimal": True,
"objective_value": self.model.ObjVal,
"partition_assignment": partition_assignment,
"topic_totals": topic_totals,
"broker_loads": broker_loads,
}
# ------------------------------------------------------------------
# Display
# ------------------------------------------------------------------
def display_results(self):
"""Print a formatted summary of the optimisation results."""
solution = self.get_solution()
if not solution["is_optimal"]:
print(f"No optimal solution found. Gurobi status: {solution['status']}")
return
print("=" * 60)
print(" Kafka Partition Assignment — Optimal Solution")
print("=" * 60)
print(f" Total inter-broker traffic cost : {solution['objective_value']:.2f}")
print()
# Partition assignment table
print(" Partition Assignment (topic x broker):")
header = (
f" {'Topic':<8}"
+ "".join(f"{b:>8}" for b in self.brokers)
+ f" {'Total':>8}"
)
print(header)
print(" " + "-" * (len(header) - 2))
for t in self.topics:
row = f" {t:<8}"
for i in self.brokers:
val = int(round(self._x[t, i].X))
row += f"{val:>8}"
row += f" {solution['topic_totals'][t]:>8}"
print(row)
print(" " + "-" * (len(header) - 2))
loads_row = f" {'Load':<8}" + "".join(
f"{solution['broker_loads'][b]:>8}" for b in self.brokers
)
print(loads_row)
print()
# Per-broker load balance
total = sum(solution["broker_loads"].values())
avg = total / len(self.brokers)
print(f" Fleet average load : {avg:.1f} partitions")
for b in self.brokers:
load = solution["broker_loads"][b]
pct = 100.0 * (load - avg) / avg if avg > 0 else 0
print(f" {b}: {load:>4} partitions ({pct:+.1f}% vs avg)")
print("=" * 60)
# ---------------------------------------------------------------------------
# Sample data
# ---------------------------------------------------------------------------
def get_sample_data():
"""Return the standard dataset from the specification."""
brokers_df = pd.DataFrame({
"broker_id": ["B1", "B2", "B3", "B4"],
"max_capacity": [50, 60, 55, 65],
})
topics_df = pd.DataFrame({
"topic_id": ["T1", "T2", "T3", "T4"],
"min_partitions": [10, 15, 8, 12],
"max_partitions": [20, 25, 18, 22],
})
# Symmetric transfer-rate matrix (higher = costlier to co-host)
rates_df = pd.DataFrame([
("B1", "B2", 5), ("B1", "B3", 3), ("B1", "B4", 4),
("B2", "B1", 5), ("B2", "B3", 4), ("B2", "B4", 7),
("B3", "B1", 3), ("B3", "B2", 4), ("B3", "B4", 2),
("B4", "B1", 4), ("B4", "B2", 7), ("B4", "B3", 2),
], columns=["from_broker", "to_broker", "transfer_rate"])
return brokers_df, topics_df, rates_df
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
if __name__ == "__main__":
brokers_df, topics_df, rates_df = get_sample_data()
with gp.Env(empty=True) as env:
env.setParam("OutputFlag", 1)
env.setParam("TimeLimit", 60)
env.start()
with KafkaPartitionOptimizer(env) as optimizer:
optimizer.set_data(brokers_df, topics_df, rates_df)
optimizer.build_model()
optimizer.solve()
optimizer.display_results()
Alongside the implementation, Modeler writes a pytest suite with one class
of tests per scenario in the specification. Each test references its
scenario identifier, and uses small hard-coded broker, topic and transfer-rate
data plus a short solver time limit so the suite runs in well under a second.
"""
Pytest suite for KafkaPartitionOptimizer.
Covers all five test scenarios defined in the specification.
"""
import time
import math
import pytest
import pandas as pd
import gurobipy as gp
from gurobipy import GRB
from kafka_modeler import KafkaPartitionOptimizer
# ---------------------------------------------------------------------------
# Shared constants
# ---------------------------------------------------------------------------
TOLERANCE = 1e-6
TIME_LIMIT = 10 # seconds per test (Gurobi TimeLimit)
# Standard transfer-rate matrix (symmetric, 4 brokers)
STANDARD_RATES = pd.DataFrame([
("B1", "B2", 5), ("B1", "B3", 3), ("B1", "B4", 4),
("B2", "B1", 5), ("B2", "B3", 4), ("B2", "B4", 7),
("B3", "B1", 3), ("B3", "B2", 4), ("B3", "B4", 2),
("B4", "B1", 4), ("B4", "B2", 7), ("B4", "B3", 2),
], columns=["from_broker", "to_broker", "transfer_rate"])
# Standard topics (same across T1, T2, T3)
STANDARD_TOPICS = pd.DataFrame({
"topic_id": ["T1", "T2", "T3", "T4"],
"min_partitions": [10, 15, 8, 12],
"max_partitions": [20, 25, 18, 22],
})
# Standard brokers
STANDARD_BROKERS = pd.DataFrame({
"broker_id": ["B1", "B2", "B3", "B4"],
"max_capacity": [50, 60, 55, 65],
})
# ---------------------------------------------------------------------------
# Helper: build and solve a model, return (solution, elapsed_seconds)
# ---------------------------------------------------------------------------
def _run(brokers_df, topics_df, rates_df, output_flag=0):
"""Create environment, build model, solve, and return solution dict."""
with gp.Env(empty=True) as env:
env.setParam("OutputFlag", output_flag)
env.setParam("TimeLimit", TIME_LIMIT)
env.start()
with KafkaPartitionOptimizer(env) as opt:
opt.set_data(brokers_df, topics_df, rates_df)
opt.build_model()
t0 = time.time()
opt.solve()
elapsed = time.time() - t0
solution = opt.get_solution()
return solution, elapsed
# ---------------------------------------------------------------------------
# Helpers for constraint verification
# ---------------------------------------------------------------------------
def check_topic_bounds(solution, topics_df):
"""All topic totals must be within [min, max]."""
for _, row in topics_df.iterrows():
t = row["topic_id"]
total = solution["topic_totals"][t]
assert total >= row["min_partitions"] - TOLERANCE, (
f"{t}: total={total} < min={row['min_partitions']}"
)
assert total <= row["max_partitions"] + TOLERANCE, (
f"{t}: total={total} > max={row['max_partitions']}"
)
def check_broker_capacity(solution, brokers_df):
"""No broker may exceed its max_capacity."""
cap = dict(zip(brokers_df["broker_id"], brokers_df["max_capacity"]))
for b, load in solution["broker_loads"].items():
assert load <= cap[b] + TOLERANCE, (
f"{b}: load={load} > capacity={cap[b]}"
)
def check_load_balance(solution):
"""Every broker load must be within 20% of the fleet average."""
loads = list(solution["broker_loads"].values())
avg = sum(loads) / len(loads)
if avg == 0:
return
for b, load in solution["broker_loads"].items():
assert load >= 0.8 * avg - TOLERANCE, (
f"{b}: load={load} < 0.8*avg={0.8*avg:.2f}"
)
assert load <= 1.2 * avg + TOLERANCE, (
f"{b}: load={load} > 1.2*avg={1.2*avg:.2f}"
)
def check_concentration(solution, topics_df):
"""No broker may hold more than 40% of any topic's partitions."""
for _, row in topics_df.iterrows():
t = row["topic_id"]
total = solution["topic_totals"][t]
if total == 0:
continue
for b in solution["broker_loads"]:
count = solution["partition_assignment"].get((t, b), 0)
assert count <= 0.4 * total + TOLERANCE, (
f"{t} on {b}: {count} > 40% of {total}"
)
# ---------------------------------------------------------------------------
# T1 - Standard Optimization
# ---------------------------------------------------------------------------
class TestT1StandardOptimization:
"""T1: Run the model with the standard 4-broker / 4-topic dataset.
Verify optimality and that all five constraints are satisfied."""
def test_t1_is_optimal(self):
"""T1: Solution must be optimal."""
solution, elapsed = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert elapsed < TIME_LIMIT, f"Solver took {elapsed:.1f}s (limit {TIME_LIMIT}s)"
assert solution["is_optimal"], "Expected an optimal solution for the standard dataset"
def test_t1_objective_positive(self):
"""T1: Objective value must be a positive finite number."""
solution, _ = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
assert solution["objective_value"] is not None
assert solution["objective_value"] > 0
assert math.isfinite(solution["objective_value"])
def test_t1_topic_partition_bounds(self):
"""T1 — C1: Every topic's total partitions in [min, max]."""
solution, _ = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_topic_bounds(solution, STANDARD_TOPICS)
def test_t1_broker_capacity(self):
"""T1 — C2: No broker exceeds its capacity."""
solution, _ = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_broker_capacity(solution, STANDARD_BROKERS)
def test_t1_load_balance(self):
"""T1 — C3: Every broker load within 20% of fleet average."""
solution, _ = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_load_balance(solution)
def test_t1_concentration_limit(self):
"""T1 — C4: No broker holds >40% of any topic's partitions."""
solution, _ = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_concentration(solution, STANDARD_TOPICS)
def test_t1_non_empty_assignment(self):
"""T1: At least one (topic, broker) pair must have partitions assigned."""
solution, _ = _run(STANDARD_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
assert len(solution["partition_assignment"]) > 0
# ---------------------------------------------------------------------------
# T2 - Tight Capacity Case
# ---------------------------------------------------------------------------
class TestT2TightCapacity:
"""T2: Broker capacities are set just barely sufficient to accommodate the
minimum partition requirements across all topics (sum min_partitions = 45).
The model should still find a feasible optimal solution."""
TIGHT_BROKERS = pd.DataFrame({
"broker_id": ["B1", "B2", "B3", "B4"],
"max_capacity": [12, 12, 11, 12],
})
def test_t2_is_optimal(self):
"""T2: Solution must be optimal under tight capacities."""
solution, elapsed = _run(self.TIGHT_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert elapsed < TIME_LIMIT
assert solution["is_optimal"], "Expected a feasible solution with tight capacities"
def test_t2_topic_partition_bounds(self):
"""T2 — C1: Topic totals within [min, max] under tight capacities."""
solution, _ = _run(self.TIGHT_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_topic_bounds(solution, STANDARD_TOPICS)
def test_t2_broker_capacity(self):
"""T2 — C2: No broker exceeds tight capacity."""
solution, _ = _run(self.TIGHT_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_broker_capacity(solution, self.TIGHT_BROKERS)
def test_t2_load_balance(self):
"""T2 — C3: Load balance satisfied under tight capacities."""
solution, _ = _run(self.TIGHT_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_load_balance(solution)
def test_t2_concentration_limit(self):
"""T2 — C4: 40% concentration limit satisfied under tight capacities."""
solution, _ = _run(self.TIGHT_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["is_optimal"]
check_concentration(solution, STANDARD_TOPICS)
# ---------------------------------------------------------------------------
# T3 - Infeasible: Capacity Too Low
# ---------------------------------------------------------------------------
class TestT3InfeasibleCapacity:
"""T3: Broker capacities are set so low (total = 20) that satisfying even the
minimum partition requirements (sum = 45) is impossible."""
INFEASIBLE_BROKERS = pd.DataFrame({
"broker_id": ["B1", "B2", "B3", "B4"],
"max_capacity": [5, 5, 5, 5],
})
def test_t3_not_optimal(self):
"""T3: Model must detect infeasibility — is_optimal False."""
solution, elapsed = _run(self.INFEASIBLE_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert elapsed < TIME_LIMIT
assert not solution["is_optimal"], (
"Expected infeasible result when total capacity < sum of min_partitions"
)
def test_t3_objective_is_none(self):
"""T3: objective_value must be None when infeasible."""
solution, _ = _run(self.INFEASIBLE_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["objective_value"] is None
def test_t3_empty_assignment(self):
"""T3: partition_assignment must be empty when infeasible."""
solution, _ = _run(self.INFEASIBLE_BROKERS, STANDARD_TOPICS, STANDARD_RATES)
assert solution["partition_assignment"] == {}
# ---------------------------------------------------------------------------
# T4 - Infeasible: Load Balance vs Concentration Conflict (2 Brokers)
# ---------------------------------------------------------------------------
class TestT4InfeasibleConcentration:
"""T4: With only 2 brokers and a 40% per-broker concentration cap, the two
brokers together can hold at most 80% of any topic's partitions, meaning
the remaining 20% cannot be placed anywhere. Combined with a minimum
partition requirement > 0, the model must be infeasible."""
TWO_BROKERS = pd.DataFrame({
"broker_id": ["B1", "B2"],
"max_capacity": [50, 50],
})
ONE_TOPIC = pd.DataFrame({
"topic_id": ["T1"],
"min_partitions": [10],
"max_partitions": [20],
})
TWO_BROKER_RATES = pd.DataFrame([
("B1", "B2", 1),
("B2", "B1", 1),
], columns=["from_broker", "to_broker", "transfer_rate"])
def test_t4_not_optimal(self):
"""T4: 40% cap with 2 brokers makes the problem infeasible — is_optimal False."""
solution, elapsed = _run(
self.TWO_BROKERS, self.ONE_TOPIC, self.TWO_BROKER_RATES
)
assert elapsed < TIME_LIMIT
assert not solution["is_optimal"], (
"Expected infeasibility: 2 brokers * 40% = 80% < 100% of partitions"
)
def test_t4_objective_is_none(self):
"""T4: objective_value must be None when infeasible."""
solution, _ = _run(self.TWO_BROKERS, self.ONE_TOPIC, self.TWO_BROKER_RATES)
assert solution["objective_value"] is None
# ---------------------------------------------------------------------------
# T5 - High Transfer Rate Avoidance
# ---------------------------------------------------------------------------
class TestT5HighTransferRateAvoidance:
"""T5: One broker pair (B1-B2) carries a very high transfer rate (1000).
The optimizer should encode this penalty correctly:
- The solution must remain feasible and optimal.
- The objective value in the high-rate scenario must be strictly greater
than the objective in an equivalent uniform-rate scenario (all rates = 1),
proving the model correctly incorporates the rate matrix into the cost.
- The B1-B2 product in the optimal solution must equal the minimum
achievable value under the active constraints (the optimizer does not
wastefully inflate it).
"""
FOUR_BROKERS = pd.DataFrame({
"broker_id": ["B1", "B2", "B3", "B4"],
"max_capacity": [50, 60, 55, 65],
})
# Single topic: with 4 brokers and min=4, load-balance forces each broker
# to hold exactly 1 partition, making product x[B1]*x[B2] = 1 unavoidable.
SINGLE_TOPIC = pd.DataFrame({
"topic_id": ["T1"],
"min_partitions": [4],
"max_partitions": [8],
})
# High rate on B1-B2 pair, all others = 1
HIGH_RATE_RATES = pd.DataFrame([
("B1", "B2", 1000), ("B1", "B3", 1), ("B1", "B4", 1),
("B2", "B1", 1000), ("B2", "B3", 1), ("B2", "B4", 1),
("B3", "B1", 1), ("B3", "B2", 1), ("B3", "B4", 1),
("B4", "B1", 1), ("B4", "B2", 1), ("B4", "B3", 1),
], columns=["from_broker", "to_broker", "transfer_rate"])
# Identical topology but uniform rate = 1 everywhere
UNIFORM_RATE_RATES = pd.DataFrame([
("B1", "B2", 1), ("B1", "B3", 1), ("B1", "B4", 1),
("B2", "B1", 1), ("B2", "B3", 1), ("B2", "B4", 1),
("B3", "B1", 1), ("B3", "B2", 1), ("B3", "B4", 1),
("B4", "B1", 1), ("B4", "B2", 1), ("B4", "B3", 1),
], columns=["from_broker", "to_broker", "transfer_rate"])
def test_t5_is_optimal(self):
"""T5: Model must find an optimal solution."""
solution, elapsed = _run(
self.FOUR_BROKERS, self.SINGLE_TOPIC, self.HIGH_RATE_RATES
)
assert elapsed < TIME_LIMIT
assert solution["is_optimal"]
def test_t5_all_constraints_satisfied(self):
"""T5: All constraints satisfied in the high-rate scenario."""
solution, _ = _run(
self.FOUR_BROKERS, self.SINGLE_TOPIC, self.HIGH_RATE_RATES
)
assert solution["is_optimal"]
check_topic_bounds(solution, self.SINGLE_TOPIC)
check_broker_capacity(solution, self.FOUR_BROKERS)
check_load_balance(solution)
check_concentration(solution, self.SINGLE_TOPIC)
def test_t5_high_rate_increases_objective(self):
"""
T5: The high-rate scenario's objective must be strictly greater than
the uniform-rate scenario's objective on the same topology.
With a single topic and 4 balanced brokers, load balance forces each
broker to hold exactly 1 partition (the only feasible assignment).
The two models therefore share the same assignment, but the high-rate
model penalises the B1-B2 pair with rate 1000 instead of 1.
Expected: obj_high > obj_uniform.
"""
solution_high, _ = _run(
self.FOUR_BROKERS, self.SINGLE_TOPIC, self.HIGH_RATE_RATES
)
solution_uniform, _ = _run(
self.FOUR_BROKERS, self.SINGLE_TOPIC, self.UNIFORM_RATE_RATES
)
assert solution_high["is_optimal"]
assert solution_uniform["is_optimal"]
assert solution_high["objective_value"] > solution_uniform["objective_value"], (
f"High-rate obj={solution_high['objective_value']:.2f} should be > "
f"uniform-rate obj={solution_uniform['objective_value']:.2f}"
)
def test_t5_minimum_possible_product(self):
"""
T5: With a single topic (min=4) and 4 brokers, the load-balance
constraint forces an exactly even distribution [1,1,1,1].
Confirm x[T1,B1] * x[T1,B2] == 1 (the minimum achievable product)
and that the optimizer does not inflate it above this floor.
"""
solution, _ = _run(
self.FOUR_BROKERS, self.SINGLE_TOPIC, self.HIGH_RATE_RATES
)
assert solution["is_optimal"]
b1 = solution["partition_assignment"].get(("T1", "B1"), 0)
b2 = solution["partition_assignment"].get(("T1", "B2"), 0)
# Load balance forces each of the 4 brokers to hold exactly 1 partition
assert b1 == 1, f"Expected x[T1,B1]=1 (forced by load balance), got {b1}"
assert b2 == 1, f"Expected x[T1,B2]=1 (forced by load balance), got {b2}"
assert b1 * b2 == 1, (
f"Minimum achievable product is 1 (load balance forces it), "
f"but got {b1}*{b2}={b1*b2}"
)
Generating code and tests is the easy half. Modeler also runs the suite in an isolated environment created fresh for this conversation, solving against Gurobi in the cloud — you do not need Gurobi installed locally, and you do not need a license of your own to see the model solve. It reads the output and reports back, so what you get back is a model that has been executed, not one that merely looks plausible.
That covers everything — please finalize the specification, then implement it in gurobipy, write a pytest suite covering the five test scenarios, run the tests, and report the results.
No open questions left to ask — every field the implementation needs
(the quadratic traffic objective, the four constraint families, the
broker and topic ranges) is already pinned down in the specification,
and this problem is small and clean enough that there is no genuinely
open structural choice left either. Writing the
KafkaPartitionOptimizer class with one integer variable per
(topic, broker) pair and NonConvex = 2 for the bilinear objective,
then a pytest suite covering T1-T5, then running both the script and
the suite.
The run, and the loop
The first full run of the suite turned up one failure. Here is the unedited output.
$ pytest test_kafka_modeler.py -v --tb=short
============================= test session starts ==============================
collected 21 items
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_is_optimal PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_objective_positive PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_topic_partition_bounds PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_broker_capacity PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_load_balance PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_concentration_limit PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_non_empty_assignment PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_is_optimal PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_topic_partition_bounds PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_broker_capacity PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_load_balance PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_concentration_limit PASSED
test_kafka_modeler.py::TestT3InfeasibleCapacity::test_t3_not_optimal PASSED
test_kafka_modeler.py::TestT3InfeasibleCapacity::test_t3_objective_is_none PASSED
test_kafka_modeler.py::TestT3InfeasibleCapacity::test_t3_empty_assignment PASSED
test_kafka_modeler.py::TestT4InfeasibleConcentration::test_t4_not_optimal PASSED
test_kafka_modeler.py::TestT4InfeasibleConcentration::test_t4_objective_is_none PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_is_optimal PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_all_constraints_satisfied PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_avoids_high_rate_pair FAILED
=================================== FAILURES ===================================
________ TestT5HighTransferRateAvoidance.test_t5_avoids_high_rate_pair _________
test_kafka_modeler.py:363: in test_t5_avoids_high_rate_pair
assert product_b1_b2 == 0, (
E AssertionError: Expected x[T1,B1]*x[T1,B2]=0 but got x[T1,B1]=1, x[T1,B2]=1
E assert 1 == 0
=========================== short test summary info ============================
FAILED test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_avoids_high_rate_pair
========================= 20 passed, 1 failed in 0.41s ==========================
This is a test error, not a modeling error. test_t5_avoids_high_rate_pair
assumed the optimizer could drive x[T1,B1]*x[T1,B2] to 0 — but T5’s own
data (a single topic with min_partitions=4 and 4 brokers) puts C3’s
load-balance rule in full force: with only 4 partitions to place across
4 brokers, ±20% of the average forces exactly one partition per broker,
so x[T1,B1]=1 and x[T1,B2]=1 in every feasible solution, independent
of the transfer rate. The scenario’s own stated intent was to confirm
the rate matrix is correctly embedded in the cost, not to confirm the
pair can be avoided entirely — the assertion was checking the wrong
thing. The fix belongs in the test, not the model.
Found it — and it is a test bug, not a model bug. test_t5_avoids_high_rate_pair
expected the optimizer to drive x[T1,B1]*x[T1,B2] to 0, but T5’s own
load-balance constraint (C3) forces an exact 1-1-1-1 split across the 4
brokers whenever a topic needs only 4 partitions — so that product is
always 1, regardless of how expensive the B1-B2 pair is. The rate
matrix is still doing its job; the test was just asserting an
unreachable outcome. I replaced it with two assertions that check what
T5 actually intends: the high-rate scenario’s objective is strictly
higher than an equivalent uniform-rate scenario on the same topology
(proving the rate matrix is embedded correctly), and the B1-B2 product
equals 1, the minimum achievable value (proving the optimizer does not
inflate it further). Re-running the full suite:
$ pytest test_kafka_modeler.py -v --tb=short
============================= test session starts ==============================
collected 21 items
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_is_optimal PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_objective_positive PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_topic_partition_bounds PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_broker_capacity PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_load_balance PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_concentration_limit PASSED
test_kafka_modeler.py::TestT1StandardOptimization::test_t1_non_empty_assignment PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_is_optimal PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_topic_partition_bounds PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_broker_capacity PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_load_balance PASSED
test_kafka_modeler.py::TestT2TightCapacity::test_t2_concentration_limit PASSED
test_kafka_modeler.py::TestT3InfeasibleCapacity::test_t3_not_optimal PASSED
test_kafka_modeler.py::TestT3InfeasibleCapacity::test_t3_objective_is_none PASSED
test_kafka_modeler.py::TestT3InfeasibleCapacity::test_t3_empty_assignment PASSED
test_kafka_modeler.py::TestT4InfeasibleConcentration::test_t4_not_optimal PASSED
test_kafka_modeler.py::TestT4InfeasibleConcentration::test_t4_objective_is_none PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_is_optimal PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_all_constraints_satisfied PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_high_rate_increases_objective PASSED
test_kafka_modeler.py::TestT5HighTransferRateAvoidance::test_t5_minimum_possible_product PASSED
============================== 21 passed in 13.94s ==============================
All 21 tests pass. Here is the solver log from the standalone run on the full dataset, plus a report of everything that was built:
Set parameter TimeLimit to value 60
Set parameter NonConvex to value 2
Gurobi Optimizer version 13.0.2 build v13.0.2rc1 (armlinux64 - "Debian GNU/Linux 13 (trixie)")
CPU model: ARM64
Thread count: 2 physical cores, 2 logical processors, using up to 2 threads
Non-default parameters:
TimeLimit 60
NonConvex 2
Optimize a model with 36 rows, 16 columns and 240 nonzeros (Min)
Model fingerprint: 0x01050349
Model has 0 linear objective coefficients
Model has 24 quadratic objective terms
Variable types: 0 continuous, 16 integer (0 binary)
Coefficient statistics:
Matrix range [4e-01, 3e+00]
Objective range [0e+00, 0e+00]
QObjective range [8e+00, 3e+01]
Bounds range [0e+00, 0e+00]
RHS range [8e+00, 6e+01]
Presolve time: 0.00s
Presolved: 68 rows, 356 columns, 936 nonzeros
Presolved model has 40 quadratic objective terms
Variable types: 0 continuous, 356 integer (340 binary)
Root relaxation: objective -2.881673e+02, 79 iterations, 0.00 seconds (0.00 work units)
Explored 4814 nodes (56497 simplex iterations) in 1.45 seconds (0.99 work units)
Solution count 10: 1260 1262 1266 ... 1470
Optimal solution found (tolerance 1.00e-04)
Best objective 1.260000000000e+03, best bound 1.260000000000e+03, gap 0.0000%
Results
Optimal plan — total traffic cost 1260, 45 partitions placed across 4 brokers:
Topic |
B1 |
B2 |
B3 |
B4 |
Total |
|---|---|---|---|---|---|
T1 |
4 |
4 |
2 |
0 |
10 |
T2 |
3 |
0 |
6 |
6 |
15 |
T3 |
3 |
3 |
2 |
0 |
8 |
T4 |
3 |
2 |
3 |
4 |
12 |
Broker |
Load |
Capacity |
vs. fleet average (11.2) |
|---|---|---|---|
B1 |
13 |
50 |
+15.6% |
B2 |
9 |
60 |
-20.0% (at the limit) |
B3 |
13 |
55 |
+15.6% |
B4 |
10 |
65 |
-11.1% |
T1 and T2 both hit the 40% per-topic concentration cap exactly (T1 on B1 and B2, T2 on B3 and B4) — the optimizer leans as hard as the rule allows onto the cheapest broker pair, B3-B4 (transfer rate 2), and away from the costliest, B2-B4 (rate 7).
Files created
File |
Description |
|---|---|
|
Input broker, topic, and transfer-rate datasets |
|
Complete problem specification |
|
Gurobipy implementation ( |
|
Pytest suite covering scenarios T1-T5 |
Test results
Test |
Scenario |
Result |
|---|---|---|
T1 |
Standard optimization — 4 brokers, 4 topics, traffic cost 1260 |
PASSED |
T2 |
Tight capacity — feasible with no slack to spare |
PASSED |
T3 |
Total broker capacity below the sum of topic minimums is infeasible |
PASSED |
T4 |
2 brokers under the 40% concentration cap is structurally infeasible |
PASSED |
T5 |
A high transfer rate on one broker pair strictly raises the objective |
PASSED |
Let me know if you would like any changes, such as a rack-awareness rule, or generating the mathematical formulation.
That log lands in the workspace as console.log alongside the code, so the
solve is auditable rather than asserted.
The loop can end one of two ways. Either every test passes, as here — or Modeler concludes a failure reflects a genuine contradiction between the specification, the implementation and the tests. In that case it stops instead of guessing, tells you what it found, and asks which side is the one to correct. That second path is what the Changing the model tab demonstrates, when a later change surfaces exactly this kind of contradiction.
What the cloud environment gives you, and what it does not
The environment is built for validating a model, not for running production
workloads. Python 3 is available with gurobipy, gurobipy-pandas,
numpy, pandas, scipy and pytest pre-installed. Within that:
Solves use a Gurobi trial license, capped at 2,000 variables and 2,000 constraints. This model, at 16 columns and 36 rows, is nowhere near the cap — but a full-scale cluster with hundreds of topics and brokers could be.
Execution is limited to 60 seconds per command, with tight CPU and memory limits. This is why the generated tests use small hard-coded data and a short solver time limit.
Additional packages cannot be installed, there is no internet access, and Compute Server and Instant Cloud cannot be used.
So treat a passing suite as evidence that the model is right, not as a performance measurement. Download the implementation and run it against your own license and your own full cluster data for that. See Solving limits for the details.
Business problems do not hold still. The specification above puts rack-awareness constraints out of scope — a reasonable simplification, until platform tells you the cluster’s brokers actually sit in physical racks, and a full-rack failure should never be able to take out most of a topic’s partitions at once. Ask for the change in plain language and Modeler updates the specification first, then propagates it to the implementation and the tests, in that order.
One more requirement: our brokers are now grouped into two racks for fault tolerance — B1 and B2 sit in Rack A, B3 and B4 sit in Rack B. For resilience, we don’t want too many of a topic’s partitions concentrated in a single rack, similar to the existing per-broker rule. Please update the specification first, then the implementation, then the tests, in that order, and tell me if anything about the current rules or test scenarios becomes ambiguous under this change.
This is a scope change, not a data correction: “rack-awareness constraints” is currently listed under Out of Scope. Updating the specification first — add a rack_id field to the broker input, a new C5 constraint for the rack cap, and remove the out-of-scope bullet. But before picking a percentage, I should check whether the obvious analogue of the existing 40% broker rule even works here: with only 2 racks, a 40% cap would make every assignment infeasible, the same structural conflict T4 already exercises at the broker level. 60% is the tightest cap that stays feasible for 2 racks. Then the code, then the tests — and T1’s baseline plan needs re-checking against the new rule, since it currently puts 80% of some topics in one rack.
That moves rack-awareness from out of scope into the model. I have updated the specification:
rack_id field to the Broker entity (Rack A: B1, B2; Rack B: B3, B4).One consequence worth flagging before I implement it: I picked 60%, not 50%. With exactly 2 racks, a 50% cap would force every topic into an exact even split, and a 40% cap (the same threshold as the existing per-broker rule) would be immediately infeasible — 2 racks × 40% = 80% of a topic’s partitions placeable, at most. 60% is the tightest limit that still guarantees the minority rack always holds at least 40% of every topic. I’ll implement 60% unless you tell me otherwise — worth confirming once you see the numbers, since the stricter 50% reading is also defensible and would change the plan.
Before you finalize — you picked 60% as the rack concentration cap, reasoning it’s the tightest feasible threshold for 2 racks. But there’s a stricter reading someone might actually want here: a strict “no rack majority” rule capping each rack at 50%, so traffic-minimization can never lean on one rack holding most of a topic’s partitions, even if 60% is technically feasible. Do the 50% and 60% readings actually produce different optimal objective values and rack loads on the baseline dataset, or does the optimizer land on the same plan either way?
Specification, then code, then tests
The specification changes first. Section 3 gains a new constraint:
C5: Per-Topic Rack Concentration Limit (60% Rule)
For any single topic, no rack may hold more than 60% of that topic's total assigned
partitions. This ensures that at least 40% of every topic's partitions always reside
in the minority rack, so a full rack failure never wipes out more than 60% of a
topic's data.
The 60% cap is the tightest limit that remains feasible for a 2-rack cluster: a 50%
cap would require exactly even splits on every integer-sized topic, while 60% allows
natural rounding while still guaranteeing meaningful cross-rack redundancy.
The implementation follows the specification, not the other way around.
set_data gains rack bookkeeping, and a new
_add_rack_concentration_limit mirrors the existing per-topic broker
concentration constraint one level up, at the rack:
def set_data(self, brokers_df, topics_df, rates_df):
self.brokers = brokers_df["broker_id"].tolist()
self.max_capacity = dict(
zip(brokers_df["broker_id"], brokers_df["max_capacity"])
)
+
+ # Rack membership
+ self.broker_rack = dict(
+ zip(brokers_df["broker_id"], brokers_df["rack_id"])
+ )
+ self.racks = brokers_df["rack_id"].unique().tolist()
+ self.rack_brokers = {
+ rack: brokers_df.loc[
+ brokers_df["rack_id"] == rack, "broker_id"
+ ].tolist()
+ for rack in self.racks
+ }
+
self.topics = topics_df["topic_id"].tolist()
...
def _add_constraints(self):
self._add_topic_partition_bounds()
self._add_broker_capacity()
self._add_load_balance()
- self._add_concentration_limit()
+ self._add_broker_concentration_limit() # renamed, still the 40% rule
+ self._add_rack_concentration_limit()
+ def _add_rack_concentration_limit(self):
+ """
+ C5 — Per-Topic Rack Concentration Limit (60% Rule):
+ No single rack may hold more than 60% of any topic's total partitions.
+
+ sum_{i in rack r} x[t, i] <= 0.6 * sum_{i'} x[t, i']
+ for all topics t, all racks r
+ """
+ for t in self.topics:
+ topic_total = gp.quicksum(self._x[t, i] for i in self.brokers)
+ for rack, rack_broker_list in self.rack_brokers.items():
+ rack_total = gp.quicksum(
+ self._x[t, i] for i in rack_broker_list
+ )
+ self.model.addConstr(
+ rack_total <= 0.60 * topic_total,
+ name=f"c5_rack_conc_{t}_{rack}"
+ )
T1, T2 and T5 each gain a rack-concentration assertion, and the suite gains a new scenario:
class TestT6InfeasibleRackConcentration:
"""T6: Both brokers are placed in Rack A (no Rack B). Any feasible assignment
places 100% of every topic's partitions in Rack A, which violates the 60%
rack concentration cap whenever min_partitions > 0."""
def test_t6_not_optimal(self):
"""T6: Single-rack cluster violates 60% cap — is_optimal must be False."""
solution, elapsed = _run(
self.ONE_RACK_BROKERS, self.ONE_TOPIC, self.TWO_BROKER_RATES
)
assert elapsed < TIME_LIMIT
assert not solution["is_optimal"], (
"Expected infeasibility: single rack holds 100% > 60% cap"
)
Modeler re-runs the whole suite, not just the new tests — 28/28 pass. The optimal objective moves from 1260 to 1322, and the plan reshuffles which brokers within each rack host which topics.
Where it stops and asks
Before settling on 60%, there was a genuine fork in how tightly to read “no rack should be over-concentrated” — and Modeler had already flagged it, rather than picking a side silently.
Yes — this is exactly the kind of fork that should be surfaced, not assumed. Solving both caps on the real dataset to check: they do not converge on the same plan. Under 60%, two topics (T1 and T2) are pinned exactly at the ceiling — the optimizer is actively using the 10-percentage-point window between 50% and 60% to lean toward the cheapest broker pair. Force that window shut at 50% and every topic is pushed to an exact even split, which costs real traffic.
Confirmed — the two readings diverge, and it isn’t a wash:
Cap |
Objective |
vs. no rack rule (1260) |
Behavior |
|---|---|---|---|
60% (implemented) |
1322 |
+62 (+4.9%) |
T1 and T2 sit exactly at the 60% ceiling; T3 and T4 land at 50/50 on their own |
50% (rejected) |
1542 |
+282 (+22.4%) |
Every topic forced to an exact 50/50 rack split, no exceptions |
Topic |
Rack A share (60% cap) |
Rack A share (50% cap) |
Binding at 60%? |
Binding at 50%? |
|---|---|---|---|---|
T1 |
60.0% |
50.0% |
Yes |
Yes |
T2 |
40.0% |
50.0% |
Yes |
Yes |
T3 |
50.0% |
50.0% |
No |
No |
T4 |
50.0% |
50.0% |
No |
No |
The 50% cap costs 220 more units (+16.6%) than the 60% cap on this exact dataset — not a wash, a real price for stricter parity. Which reading matches your actual redundancy requirement?
Let’s go with the 60% cap — it matches what’s already implemented, and “guard against severe imbalance without forcing exact parity” is the right policy for us; we care about traffic cost, and 16.6% is a real price to pay for symmetry we don’t strictly need. Keep what’s implemented.
This is the same “stop and ask” behavior in a different shape: not a test scenario whose expected answer became ambiguous, but a modeling choice inside the new constraint that had two defensible readings with different numeric consequences. Modeler surfaced both, with real numbers attached, rather than silently keeping the one it had already coded.
The plan reshuffles, the broker loads do not
Topic |
Broker split (no rack rule) |
Broker split (60% rack rule) |
Change |
|---|---|---|---|
T1 |
B1=4, B2=4, B3=2, B4=0 |
B1=2, B2=4, B3=4, B4=0 |
Rack A share drops from 80% (would violate C5) to 60% — 2 partitions move from B1 to B3 |
T2 |
B1=3, B2=0, B3=6, B4=6 |
B1=6, B2=0, B3=3, B4=6 |
Rack B share drops from 80% (would violate C5) to 60% — 3 partitions move from B3 to B1 |
T3 |
B1=3, B2=3, B3=2, B4=0 |
B1=1, B2=3, B3=2, B4=2 |
Rack A share drops from 75% to an even 50% — 2 partitions move from B1 to B4 |
T4 |
B1=3, B2=2, B3=3, B4=4 |
B1=4, B2=2, B3=4, B4=2 |
Already at 58.3% Rack B, coincidentally under 60% — reshuffles anyway to free room elsewhere |
Broker loads (B1/B2/B3/B4) |
13 / 9 / 13 / 10 |
13 / 9 / 13 / 10 |
Unchanged — C3’s load-balance point stays put; only which topics land where, within each rack, shifts |
Objective |
1260 |
1322 |
+62 (+4.9%) — the real cost of guaranteeing every topic survives a full-rack failure with at least 40% intact |
Rejected 50% reading |
— |
1542 |
Would have cost +282 (+22.4%) instead |
Adding the rack constraint does not touch how much load each broker carries — C3’s 20% load-balance rule pins that independently. What changes is which topics each broker within a rack hosts: T1 and T2, the two topics that were leaning hardest on the cheap B3-B4 pair, both get pulled back to a legal 60/40 split, and the small ripple from that reshuffling 2-3 partitions at a time is the entire +62 difference between the old and new optimal plans.
The specification, the implementation and the tests never mention mathematics — that is deliberate, so the business owner can read them. But you may still need the formulation: to review the model, to document it, or to hand it to someone who thinks in constraints rather than in code. Ask at any point and Modeler derives one from the current specification and implementation, in Markdown with LaTeX, and saves it to the workspace with everything else.
Sets and indices
\(t \in T\): topics.
\(i, j \in B\): brokers.
\(P = \{(i,j) \in B \times B \mid i \neq j\}\): ordered broker pairs.
Parameters
\(\underline{p}_t \in \mathbb{Z}_{\geq 0}\): minimum number of partitions required for topic \(t \in T\).
\(\overline{p}_t \in \mathbb{Z}_{\geq 0}\): maximum number of partitions allowed for topic \(t \in T\).
\(c_i \in \mathbb{Z}_{\geq 0}\): maximum partition capacity of broker \(i \in B\).
\(r_{ij} \in \mathbb{R}_{\geq 0}\): pairwise inter-broker transfer rate for ordered pair \((i,j) \in P\).
\(\alpha = 0.20\): load-balance tolerance (20% rule).
\(\beta = 0.40\): per-topic broker concentration limit (40% rule).
The transfer-rate matrix is symmetric: \(r_{ij} = r_{ji}\) for all \((i,j) \in P\). Both directions are kept in the formulation so the objective sums each unordered pair \(\{i,j\}\) twice, exactly consistent with the implementation.
Decision variables
\(x_{ti} \in \mathbb{Z}_{\geq 0}\) for \(t \in T,\ i \in B\): the integer number of partitions of topic \(t\) assigned to broker \(i\).
Objective function
Minimize total inter-broker traffic cost, the sum over all topics and all ordered broker pairs of the transfer rate weighted by the product of partition counts on the two brokers:
The products \(x_{ti} \cdot x_{tj}\) are bilinear terms in integer variables. Because \(r_{ij} = r_{ji}\), the double sum over all ordered pairs counts each unordered pair \(\{i,j\}\) exactly twice; the factor of 2 scales the objective uniformly and does not alter the optimal solution.
Constraints
C1 — Topic partition bounds.
C2 — Broker capacity.
C3 — Global load balance (20% rule). Letting \(X = \sum_{t \in T}\sum_{i \in B} x_{ti}\) be the total partitions placed, every broker’s load must lie within \(\pm\alpha\) of the fleet average \(X / |B|\). Multiplying through by \(|B|\) keeps both sides linear in the decision variables:
C4 — Per-topic broker concentration limit (40% rule).
C5 — Non-negativity and integrality.
Complete model
Additional notes
The bilinear products \(x_{ti} \cdot x_{tj}\) of integer variables are neither convex nor concave in general, so this is a non-convex MIQP. Gurobi requires
NonConvex = 2to certify global optimality via spatial branch-and-bound, which internally linearizes the bilinear terms.On the 4-broker / 4-topic baseline instance, the model has \(4 \times 4 = 16\) integer decision variables and 36 linear constraint rows before presolve — matching the solver log exactly (
Optimize a model with 36 rows, 16 columns and 240 nonzeros ... Model has 24 quadratic objective terms; 24 rather than 48 because Gurobi stores each symmetric pair once internally).C4’s per-topic concentration limit is binding for two topics on the baseline dataset (T1 on B1 and B2 at exactly 40%; T2 on B3 and B4 at exactly 40%) — the optimizer uses the full permitted concentration to lean toward the cheapest broker pair, B3-B4 (rate 2), and away from the costliest, B2-B4 (rate 7).
Note that this formulates the model before the rack-awareness requirement introduced in the Changing the model tab. Add constraint C5’ — \(\sum_{i \in \text{rack}(r)} x_{ti} \leq 0.60 \cdot \sum_{j \in B} x_{tj}\) for every topic \(t\) and rack \(r\) — and renumber the non-negativity / integrality constraint to C6, to recover the version formulated after that change.
### Sets and indices
- $t \in T$: topics.
- $i, j \in B$: brokers.
- $P = \{(i,j) \in B \times B \mid i \neq j\}$: ordered broker pairs.
### Parameters
- $\underline{p}_t \in \mathbb{Z}_{\geq 0}$: minimum number of partitions required for topic $t \in T$.
- $\overline{p}_t \in \mathbb{Z}_{\geq 0}$: maximum number of partitions allowed for topic $t \in T$.
- $c_i \in \mathbb{Z}_{\geq 0}$: maximum partition capacity of broker $i \in B$.
- $r_{ij} \in \mathbb{R}_{\geq 0}$: pairwise inter-broker transfer rate for ordered pair $(i,j) \in P$.
- $\alpha = 0.20$: load-balance tolerance (20% rule).
- $\beta = 0.40$: per-topic broker concentration limit (40% rule).
The transfer-rate matrix is symmetric: $r_{ij} = r_{ji}$ for all $(i,j) \in P$.
Both directions are kept in the formulation so the objective sums each
unordered pair $\{i,j\}$ twice, exactly consistent with the implementation.
### Decision variables
- $x_{ti} \in \mathbb{Z}_{\geq 0}$ for $t \in T,\ i \in B$: the integer number of partitions of topic $t$ assigned to broker $i$.
### Objective function
Minimize total inter-broker traffic cost, the sum over all topics and all
ordered broker pairs of the transfer rate weighted by the product of
partition counts on the two brokers:
$$
\min \quad \sum_{t \in T} \sum_{\substack{i \in B,\ j \in B \\ i \neq j}} r_{ij}\, x_{ti}\, x_{tj}
$$
The products $x_{ti} \cdot x_{tj}$ are bilinear terms in integer variables.
Because $r_{ij} = r_{ji}$, the double sum over all ordered pairs counts each
unordered pair $\{i,j\}$ exactly twice; the factor of 2 scales the objective
uniformly and does not alter the optimal solution.
### Constraints
**C1 — Topic partition bounds.**
$$
\sum_{i \in B} x_{ti} \geq \underline{p}_t \qquad \forall\, t \in T
$$
$$
\sum_{i \in B} x_{ti} \leq \overline{p}_t \qquad \forall\, t \in T
$$
**C2 — Broker capacity.**
$$
\sum_{t \in T} x_{ti} \leq c_i \qquad \forall\, i \in B
$$
**C3 — Global load balance (20% rule).** Letting $X = \sum_{t \in T}\sum_{i \in B} x_{ti}$
be the total partitions placed, every broker's load must lie within $\pm\alpha$
of the fleet average $X / |B|$. Multiplying through by $|B|$ keeps both sides
linear in the decision variables:
$$
|B| \cdot \sum_{t \in T} x_{ti} \geq (1-\alpha)\, X \qquad \forall\, i \in B
$$
$$
|B| \cdot \sum_{t \in T} x_{ti} \leq (1+\alpha)\, X \qquad \forall\, i \in B
$$
**C4 — Per-topic broker concentration limit (40% rule).**
$$
x_{ti} \leq \beta \cdot \sum_{j \in B} x_{tj} \qquad \forall\, t \in T,\ i \in B
$$
**C5 — Non-negativity and integrality.**
$$
x_{ti} \in \mathbb{Z}_{\geq 0} \qquad \forall\, t \in T,\ i \in B
$$
### Complete model
$$
\begin{aligned}
\min \quad & \sum_{t \in T} \sum_{\substack{i \in B,\ j \in B \\ i \neq j}} r_{ij}\, x_{ti}\, x_{tj} \\[6pt]
\text{s.t.} \quad & \sum_{i \in B} x_{ti} \geq \underline{p}_t && \forall\, t \in T & \text{(C1a)} \\[4pt]
& \sum_{i \in B} x_{ti} \leq \overline{p}_t && \forall\, t \in T & \text{(C1b)} \\[4pt]
& \sum_{t \in T} x_{ti} \leq c_i && \forall\, i \in B & \text{(C2)} \\[4pt]
& |B| \cdot \sum_{t \in T} x_{ti} \geq (1-\alpha)\, X && \forall\, i \in B & \text{(C3a)} \\[4pt]
& |B| \cdot \sum_{t \in T} x_{ti} \leq (1+\alpha)\, X && \forall\, i \in B & \text{(C3b)} \\[4pt]
& x_{ti} \leq \beta \cdot \sum_{j \in B} x_{tj} && \forall\, t \in T,\ i \in B & \text{(C4)} \\[4pt]
& x_{ti} \in \mathbb{Z}_{\geq 0} && \forall\, t \in T,\ i \in B & \text{(C5)}
\end{aligned}
$$
### Additional notes
- The bilinear products $x_{ti} \cdot x_{tj}$ of integer variables are
neither convex nor concave in general, so this is a **non-convex MIQP**.
Gurobi requires `NonConvex = 2` to certify global optimality via spatial
branch-and-bound, which internally linearizes the bilinear terms.
- On the 4-broker / 4-topic baseline instance, the model has $4 \times 4 = 16$
integer decision variables and 36 linear constraint rows before presolve —
matching the solver log exactly (`Optimize a model with 36 rows, 16 columns
and 240 nonzeros ... Model has 24 quadratic objective terms`; 24 rather
than 48 because Gurobi stores each symmetric pair once internally).
- C4's per-topic concentration limit is binding for two topics on the
baseline dataset (T1 on B1 and B2 at exactly 40%; T2 on B3 and B4 at
exactly 40%) — the optimizer uses the full permitted concentration to lean
toward the cheapest broker pair, B3-B4 (rate 2), and away from the
costliest, B2-B4 (rate 7).
Note that this formulates the model **before** the rack-awareness
requirement introduced in the Changing the model tab. Add constraint
C5' — $\sum_{i \in \text{rack}(r)} x_{ti} \leq 0.60 \cdot \sum_{j \in B} x_{tj}$
for every topic $t$ and rack $r$ — and renumber the non-negativity /
integrality constraint to C6, to recover the version formulated after that
change.