24 lines
944 B
Python
24 lines
944 B
Python
#!/usr/bin/env python3
|
|
"""Demo: ADMM-like optimization with two local problems"""
|
|
from catopt_flow_category_theoretic_compositi import LocalProblem, DeltaSyncRegistry, ADMMNode, run_admm, Planner
|
|
|
|
|
|
def main():
|
|
lp1 = LocalProblem(id="node-1", resources={"gpu": 2.0, "memory": 8.0}, data_loading={}, batch_size=32)
|
|
lp2 = LocalProblem(id="node-2", resources={"gpu": 1.0, "memory": 8.0}, data_loading={}, batch_size=32)
|
|
registry = DeltaSyncRegistry()
|
|
n1 = ADMMNode("node-1", lp1, registry)
|
|
n2 = ADMMNode("node-2", lp2, registry)
|
|
final = run_admm([n1, n2], rounds=3, budget_gpu=4.0, registry=registry)
|
|
print("Final global plan:")
|
|
for lid, alloc in final.plan.items():
|
|
print(f" {lid}: {alloc}")
|
|
|
|
# Show a planner-only view from local problems (example of usage)
|
|
plan = Planner.build_global_plan([lp1, lp2], budget_gpu=4.0)
|
|
print("Planner plan:", plan.plan)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|