In this tutorial, we implement the Rebalancer, Meta’s assignment solver, with a C++ core and a Python binding. We describe objects, containers, dimensions, scopes, and partitions, add constraints and goals, then choose local search or the exact mixed-integer solver HiGHS. We build a small datacenter, drain a busy host, and compare the two solvers on toy models, mutually exclusive hosts, synthetic clusters, and 60 published OR-Library instances. We judge every plan by recounting the returned assignment rather than by the solver’s flags. In this build, across 24 solves, the solved flag read True in 14 solves that still left memory over capacity. HARD under local search returned a violating start without raising; misspelled spec keys were ignored silently, and enum names are rejected. The timings come from an Apple M4 with the macOS wheel, and Colab’s Linux wheel may time differently.
import os
import re
import sys
import json
import math
import time
import random
import hashlib
import platform
import tempfile
import statistics
import traceback
import subprocess
import collections
import urllib.request
RESULTS = {} # step -> headline (or the failure), saved to results.json after every step
CLAIMS = [] # (step, claim, held) rows for the ledger in the last cell
SOLVE_SECONDS = collections.Counter() # step -> seconds spent inside solve()
CURRENT = ["1"]
def banner(title):
print("\n" + "=" * 78)
print(title)
print("=" * 78)
def section(name):
def wrap(fn):
def run(*a, **kw):
CURRENT[0] = name.split(".")[0]
banner(name)
t0 = time.time()
try:
out = fn(*a, **kw)
RESULTS[name] = out if isinstance(out, str) else "ok"
return out
except Exception as e:
RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
traceback.print_exc(limit=3)
return None
finally:
print(f"\n[{time.time() - t0:.1f} s]")
json.dump({"results": RESULTS, "claims": CLAIMS, "solve_seconds": SOLVE_SECONDS},
open("results.json", "w"), indent=1)
return run
return wrap
def claim(text, held):
CLAIMS.append((CURRENT[0], text, bool(held)))
return held
banner("1. Install rebalancer 1.0.4, check the build, and solve one small model with both solvers")
assert (3, 12) <= sys.version_info[:2] <= (3, 14), "the 1.0.4 wheels exist for Python 3.12, 3.13 and 3.14 only"
assert (sys.platform.startswith("linux") and platform.machine() == "x86_64") or \
(sys.platform == "darwin" and platform.machine() == "arm64"), f"no rebalancer 1.0.4 wheel for {sys.platform} {platform.machine()}"
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--retries", "5", "--timeout", "120",
"rebalancer==1.0.4"], check=True)
import importlib.metadata as md
import rebalancer.specs as S
from rebalancer import ProblemSolver
# The binding takes every enum as its Thrift integer. The type stubs advertise the names ("HARD", "HIGHS"), but a name
# raises '"" is not a valid int', and an out-of-range integer is accepted silently, so the integers live in one table.
POLICY = {"DEFAULT": 1, "HARD": 2, "SOFT": 3}
LIMIT = {"ABSOLUTE": 1, "RELATIVE": 2}
FORMULA = {"LINEAR": 1, "SQUARES": 2, "MAX": 3, "IDEAL": 4}
PACKAGE = {"XPRESS": 0, "GUROBI": 1, "HIGHS": 2}
MAX_BOUND = 1
END = {1: "OPTIMAL", 2: "HIT_TIME_LIMIT", 3: "HIT_MOVE_LIMIT", 4: "HIT_STOP_CONSTRAINT", 5: "HIT_PLATEAU_TIME",
6: "NO_SOLUTION_FOUND", 7: "HIT_STAGE_MOVE_LIMIT", 8: "HIT_EXCEPTION", 9: "UNABLE_TO_FIND_MORE_MOVES",
10: "HIT_MIN_CYCLE_OBJECTIVE_IMPROVEMENT", 11: "UNABLE_TO_FIND_IMPROVING_MOVES"}
NESTED = {"limit": "Limit", "conflictInfoList": "ScopeItemConflictInfo",
"conflictingScopeItemsWithOverlap": "ConflictingScopeItemInfo"}
def check_keys(name, fields, where):
cls = getattr(S, name)
unknown = set(fields) - (cls.__required_keys__ | cls.__optional_keys__)
if unknown:
raise KeyError(f"{where} has no field {sorted(unknown)}; the binding would ignore it silently")
for key, value in fields.items():
sub = "AssignmentAffinity" if (where, key) == ("assignmentAffinitiesSpec", "affinities") else NESTED.get(key)
for item in (value if isinstance(value, list) else [value]):
if sub and isinstance(item, dict):
check_keys(sub, item, f"{where}.{key}")
def spec(arm, **fields):
"""One union arm, e.g. spec("capacitySpec", name=..., scope=..., dimension=...).
The binding ignores unknown keys without a warning, so the keys are checked against the TypedDicts in rebalancer.specs:
the arm's own keys and the nested limit, affinities and conflict lists used in this notebook."""
check_keys(arm[0].upper() + arm[1:], fields, arm)
return {arm: fields}
MOVES = {"single": {"singleMoveTypeSpec": {}}, "swap": {"swapMoveTypeSpec": {}},
"triple": {"tripleLoopMoveTypeSpec": {}}, "kl": {"klSearchMoveTypeSpec": {}}}
def local_search(seconds=2, seed=0, moves=("single", "swap"), **extra):
return spec("localSearchSolverSpec", moveTypeList=[MOVES[m] for m in moves], solveTime=int(seconds),
randomSeed=seed, **extra)
def highs(seconds=5):
# solverPackage 2 is HiGHS, the only MIP solver compiled into the pip wheel. suppressLogs removes a banner that the
# optimal path otherwise prints to stdout.
return spec("optimalSolverSpec", solverPackage=PACKAGE["HIGHS"], solveTime=int(seconds), suppressLogs=True)
def new_problem(assignment, objects="task", containers="host"):
p = ProblemSolver(service_name="rebalancer_tutorial", service_scope="marktechpost")
p.set_log_level("ERROR") # silences ~50 log lines per solve; never "OFF", which exits the process
p.set_object_name(objects).set_container_name(containers) # names first, then the assignment, then dimensions
p.set_assignment(assignment) # every container must be listed, empty ones too
return p
def solve(p):
t = time.perf_counter()
sol = p.solve()
dt = time.perf_counter() - t
SOLVE_SECONDS[CURRENT[0]] += dt
return sol, dt
# Independent checks: everything below reads only the returned assignment, never the solver's own verdict.
def where(container_to_objects):
return {o: c for c, objs in container_to_objects.items() for o in objs}
def by_container(object_to_container, containers):
out = {c: [] for c in containers}
for o, c in sorted(object_to_container.items()):
out[c].append(o)
return out
def load(object_to_container, values):
out = collections.Counter()
for o, c in object_to_container.items():
out[c] += values.get(o, 0.0)
return out
def excess(object_to_container, values, caps):
L = load(object_to_container, values)
return {c: max(0.0, L[c] - caps[c]) for c in caps}
def moved(before, after):
return sum(before[o] != after[o] for o in before)
def ends(sol):
return "/".join(END.get(r["endReason"], str(r["endReason"])) for r in sol["solverSummaries"])
def goals(sol):
return [round(g["value"], 6) for g in sol["finalGlobalObjective"]["goals"]]
def clean(e, n=90, typed=True):
msg = re.sub(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "<id>", str(e)).strip()[:n]
return f"{type(e).__name__}: {msg}" if typed else msg
def cpu_name():
if os.path.exists("/proc/cpuinfo"):
for line in open("/proc/cpuinfo"):
if line.startswith("model name"):
return line.split(":", 1)[1].strip()
if sys.platform == "darwin":
r = subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True)
if r.stdout.strip():
return r.stdout.strip()
return platform.processor() or platform.machine()
def ram_gb():
if os.path.exists("/proc/meminfo"):
kb = next(int(l.split()[1]) for l in open("/proc/meminfo") if l.startswith("MemTotal"))
return kb / 1024 ** 2
r = subprocess.run(["sysctl", "-n", "hw.memsize"], capture_output=True, text=True)
return int(r.stdout) / 1024 ** 3 if r.stdout.strip() else float("nan")
VERSION = md.version("rebalancer")
HW = f"{cpu_name()}, {os.cpu_count()} logical CPUs, {ram_gb():.0f} GB RAM"
print(f"rebalancer {VERSION} python {platform.python_version()} {platform.system()} {platform.machine()}")
print(f"hardware: {HW}")
assert VERSION == "1.0.4", "this notebook is written against rebalancer 1.0.4; restart the runtime after installing"
def two_tasks(solver_spec, policy=None):
# Two 10 GB tasks on host0, two hosts of 15 GB: the start breaks capacity and one task has to move.
p = new_problem({"host0": ["t0", "t1"], "host1": []})
p.add_object_dimension("memory", {"t0": 10.0, "t1": 10.0})
p.add_container_dimension("memory", {}, default_value=15.0)
kw = {} if policy is None else {"policy": policy}
p.add_constraint(spec("capacitySpec", name="memory_capacity", scope="host", dimension="memory"), **kw)
p.add_solver(solver_spec)
return p
GATE = {}
for label, make in [("local search", lambda: two_tasks(local_search(2))),
("HiGHS (solverPackage=2)", lambda: two_tasks(highs(2)))]:
try:
sol, dt = solve(make())
per_host = load(sol["assignment"], {"t0": 10.0, "t1": 10.0})
GATE[label] = sol
print(f"{label:<28} -> host loads {dict(sorted(per_host.items()))} GB, endReason {ends(sol)}, "
f"{dt * 1000:.1f} ms")
except Exception as e:
GATE[label] = None
print(f"{label:<28} -> FAILED {clean(e)}")
print("\nWhat the build refuses (each is caught here; the message is the binding's own):")
for label, make in [("OptimalSolverSpec() default", lambda: two_tasks(spec("optimalSolverSpec", suppressLogs=True))),
("solverPackage=1 (Gurobi)", lambda: two_tasks(spec("optimalSolverSpec", solverPackage=1,
suppressLogs=True))),
('policy="HARD" (a name)', lambda: two_tasks(local_search(2), policy="HARD")),
('solverPackage="HIGHS"', lambda: two_tasks({"optimalSolverSpec": {"solverPackage": "HIGHS"}}))]:
try:
solve(make())
print(f" {label:<28} -> accepted")
except Exception as e:
print(f" {label:<28} -> {clean(e, 60)}")
try:
spec("localSearchSolverSpec", timeLimitMs=500)
except KeyError as e:
print(f" spec(..., timeLimitMs=500) -> {e.args[0]}")
def gate_ok(sol): # solved means OPTIMAL and, by recount, no host above its 15 GB
return (sol is not None and ends(sol) == "OPTIMAL" and len(sol["assignment"]) == 2
and max(load(sol["assignment"], {"t0": 10.0, "t1": 10.0}).values()) <= 15.0)
HIGHS_OK = gate_ok(GATE["HiGHS (solverPackage=2)"])
if not gate_ok(GATE["local search"]):
raise RuntimeError("local search did not solve the two-task model; nothing below can run on this runtime")
print(f"\nGATE: local search solved; HiGHS {'solved' if HIGHS_OK else 'FAILED (steps that need it will say so)'}")
RESULTS["1. Install, check the build, solve one small model with both solvers"] = f"rebalancer {VERSION}, Python {platform.python_version()}, {HW}; HiGHS {'ok' if HIGHS_OK else 'failed'}"
We install rebalancer 1.0.4 and print the Python version and hardware, because every timing depends on them. The binding rejects enum names and silently ignores unknown spec keys, so the notebook passes integers from one table and checks each spec before solving. A gate then solves a two-task model with both solvers, each moving one task, in 0.8 ms with local search and 0.9 ms with HiGHS. The cell also confirms that this build refuses the default XPRESS solver and the Gurobi package.
README_TASKS = {f"task{i}": 10.0 for i in range(4)}
README_CAPS = {"host0": 20.0, "host1": 20.0}
def readme_model(start, solver_spec):
p = new_problem(start)
p.add_object_dimension("memory", README_TASKS)
p.add_container_dimension("memory", {}, default_value=20.0)
# No limit given: the capacity limit is RELATIVE 1.0, i.e. 100% of each host's memory dimension.
p.add_constraint(spec("capacitySpec", name="memory_capacity", scope="host", dimension="memory"))
p.add_solver(solver_spec)
return p
@section("2. The README model: four 10 GB tasks, two 20 GB hosts, and a recount instead of the solver's verdict")
def readme_step():
start = {"host0": ["task0", "task1", "task2"], "host1": ["task3"]}
sol, dt = solve(readme_model(start, local_search(2)))
before, after = where(start), sol["assignment"]
ex0, ex1 = excess(before, README_TASKS, README_CAPS), excess(after, README_TASKS, README_CAPS)
l0, l1 = load(before, README_TASKS), load(after, README_TASKS)
ic, fc = sol["initialConstraint"], sol["finalConstraint"]
print(f"start: {start}")
print(f"final: {by_container(after, README_CAPS)}")
print(f"host loads {l0['host0']:.0f}/{l0['host1']:.0f} GB -> {l1['host0']:.0f}/{l1['host1']:.0f} GB "
f"against 20 GB; recounted excess {sum(ex0.values()):.0f} -> {sum(ex1.values()):.0f} GB")
print(f"solver's own report: brokenCount {ic['brokenCount']} -> {fc['brokenCount']}, brokenVal "
f"{ic['brokenVal']:g} -> {fc['brokenVal']:g}, solved={fc['solved']}")
print(f"tasks moved (diff against the start): {moved(before, after)} movesSummary entries: "
f"{len(sol['movesSummary'])} endReason {ends(sol)} solve {dt * 1000:.1f} ms")
print(f"result keys: {', '.join(sorted(sol))}")
balanced = {"host0": ["task0", "task1"], "host1": ["task2", "task3"]}
sol_b, _ = solve(readme_model(balanced, local_search(2)))
sol_z, _ = solve(readme_model(start, local_search(0)))
mb, mz = moved(where(balanced), sol_b["assignment"]), moved(before, sol_z["assignment"])
exz = sum(excess(sol_z["assignment"], README_TASKS, README_CAPS).values())
print(f"\ncontrol, start already 2-2: {mb} tasks moved, endReason {ends(sol_b)}")
print(f"control, solveTime=0 on the 3-1 start: {mz} tasks moved, endReason {ends(sol_z)}, recounted excess "
f"{exz:.0f} GB, solved={sol_z['finalConstraint']['solved']} <- 0 means no time, not 'no limit'")
claim("README model: recounted excess reaches 0 GB", sum(ex1.values()) == 0)
claim("README model: a 2-2 start is left alone", mb == 0)
claim("README model: solveTime=0 makes no moves", mz == 0)
return (f"3-1 start -> {l1['host0']:.0f}/{l1['host1']:.0f} GB, excess {sum(ex0.values()):.0f} -> "
f"{sum(ex1.values()):.0f} GB, {moved(before, after)} task moved; solveTime=0 moved {mz}")
readme_step()
We solve the README example, four 10 GB tasks on two 20 GB hosts with three tasks on one host, and recount the result instead of trusting the solver. Local search moves one task in 0.3 ms, and the recount drops the excess from 10 GB to 0. The solver’s own brokenCount also falls from 1 to 0 and endReason reads OPTIMAL. A control makes the point about the flag: with solveTime=0 the solver moves nothing, still leaves 10 GB over capacity, and reports solved=True.
HOSTS = [f"host{i}" for i in range(8)]
RACK = {h: f"rack{i // 2}" for i, h in enumerate(HOSTS)} # Python add_scope maps container -> scope item
JOB_OF = {f"job{j}_t{k}": f"job{j}" for j in "ABCD" for k in range(4)}
TASK_CPU, TASK_MEM, HOST_CPU, HOST_MEM = 4.0, 16.0, 16.0, 64.0
DC_START = {h: [] for h in HOSTS}
DC_START["host0"] = [t for t, j in JOB_OF.items() if j in ("jobA", "jobB")]
DC_START["host1"] = [t for t, j in JOB_OF.items() if j == "jobC"]
DC_START["host2"] = [t for t, j in JOB_OF.items() if j == "jobD"]
def datacenter(solver_spec, start=DC_START, rack_rule=True, formula="LINEAR", drain=None, movement_tuple=None,
limit_type="ABSOLUTE"):
p = new_problem(start)
p.add_partition("job", {j: [t for t in JOB_OF if JOB_OF[t] == j] for j in sorted(set(JOB_OF.values()))})
p.add_scope("rack", RACK)
p.add_object_dimension("cpu", {t: TASK_CPU for t in JOB_OF})
p.add_object_dimension("mem", {t: TASK_MEM for t in JOB_OF})
p.add_container_dimension("cpu", {}, default_value=HOST_CPU)
p.add_container_dimension("mem", {}, default_value=HOST_MEM)
p.add_constraint(spec("capacitySpec", name="cpu_capacity", scope="host", dimension="cpu", bound=MAX_BOUND))
p.add_constraint(spec("capacitySpec", name="mem_capacity", scope="host", dimension="mem", bound=MAX_BOUND))
if rack_rule: # at most one task of each job per rack, as an ABSOLUTE limit (step 3 shows what RELATIVE does)
p.add_constraint(spec("groupCountSpec", name="one_task_per_job_per_rack", scope="rack", partitionName="job",
bound=MAX_BOUND, limit={"type": LIMIT[limit_type], "globalLimit": 1.0}))
if drain:
p.add_constraint(spec("toFreeSpec", name=f"drain_{drain}", containers=[drain]))
p.add_goal(spec("balanceSpec", name="cpu_balance", scope="host", dimension="cpu", formula=FORMULA[formula]),
weight=1.0, tuple_pos=0)
if movement_tuple is not None:
p.add_goal(spec("minimizeMovementSpec", name="tasks_moved", scope="host", dimension="task_count",
doNotNormalize=True), weight=1.0, tuple_pos=movement_tuple)
p.add_solver(solver