Cluster configuration & partition reassignment

5 h
requires KFK-02

From your world. Redshift WLM queues / slot rebalancing: moving work to even out load without a full reload

Minimal-movement partition reassignment (+ apply it)

To rebalance a Kafka cluster, you recompute the replica list of each partition. But every replica that you move is a full data copy over the network. The skill is to rebalance with the fewest moves and keep every assignment valid. Edit reassign.py and implement three functions.

1. `plan_reassignment(current, broker_set, rf, racks=None) -> dict[int, list[int]]` (pure)

current: {partition: [broker, ...]} — the existing replica lists.
broker_set: the brokers where you may place replicas.
rf: target replication factor.
racks: optional {broker: rack}.

Return {partition: [broker, ...]} so that every partition meets these rules:

It has exactly rf replicas. All replicas are distinct, and all come from broker_set.
When racks is given and rf <= number_of_distinct_racks, every replica is on a distinct rack.
It keeps as many current replicas as possible. Keep a replica that is already on a broker in broker_set (and does not violate the rack rule). Keeping it costs zero data movement. Replace only the replicas that you must replace, and spread the new ones to balance the load.

2. `to_assignment_json(plan, topic) -> dict` (pure)

Return the document that kafka-reassign-partitions.sh --execute expects:

{"version": 1, "partitions": [{"topic": "t", "partition": 0, "replicas": [1, 3]}, ...]}

3. `apply_reassignment(cluster, topic, plan, timeout=90.0) -> bool`

Write the plan JSON to a temp file. Run kafka-reassign-partitions.sh ... --execute with cluster.run_tool(...). Then poll --verify until it reports that the reassignment is complete. Return True when it completes.

How it is graded

Pure validity + minimality (no broker): the grader gives random partitions, broker sets, rf, and racks to your planner (property-based). It asserts that each result is valid (rf, distinct, subset, rack spread). It also asserts that you never dropped a reusable current replica while you added a brand-new broker (minimal movement). An independent oracle checks the movement bound.
Live apply (native 3-broker cluster, no Docker): the grader creates a topic and asks your planner to move it onto a smaller broker set. It applies the plan. Then it asserts that the replica placement of the real cluster matches your plan. It also asserts that the dropped broker hosts none of the replicas of the topic.

Starter code

reassign.py

"""Starter — KFK-11. Implement the minimal-movement reassignment planner and the
apply helpers. The stubs below currently fail the tests (gate G1)."""

from __future__ import annotations

import json
import tempfile
import time
from pathlib import Path


def plan_reassignment(current, broker_set, rf, racks=None):
    """Return {partition: [broker, ...]} with exactly rf distinct replicas from
    broker_set, rack-distinct when racks are given and rf <= number of racks,
    keeping as many current replicas as possible (minimal data movement)."""
    # TODO
    return {}


def to_assignment_json(plan, topic):
    """Return {'version':1,'partitions':[{'topic','partition','replicas'}, ...]}."""
    # TODO
    return {}


def apply_reassignment(cluster, topic, plan, timeout=90.0):
    """Write the plan JSON, run kafka-reassign-partitions.sh --execute, then poll
    --verify until the reassignment is complete. Return True on completion."""
    # TODO
    return False

Hints

Hint 1

plan_reassignment works per partition. Keep every current replica that is still in the target broker set AND does not violate rack-distinctness; only those kept replicas cost zero data movement. Count how many slots you still need (rf minus kept).

Hint 2

Fill the remaining slots from brokers NOT already chosen for that partition, preferring brokers on a rack not yet used by this partition (when racks are given), and spreading load (prefer the least-used broker across all partitions). Never place two replicas on the same broker; never exceed rf.

Hint 3

to_assignment_json(plan, topic) builds the {'version':1,'partitions':[...]} dict that kafka-reassign-partitions.sh --execute expects: one entry per partition with 'topic','partition','replicas'. apply_reassignment writes it to a temp file and runs the tool via cluster.run_tool, then polls --verify until complete.

Run it locally

uv run de start KFK-11/01-minimal-reassignment
# edit workspace/KFK-11/01-minimal-reassignment/
uv run de test