Datacenter Compute Allocation
Allocate GPU capacity across hyperscaler campuses using a graph neural network (GNN) that predicts per-workload utilization, hardware-compatibility rules, dependency PageRank, and a 24-cell scenario mixed-integer program (MIP). Picks up where energy_grid_planning leaves off.
Browse files
Directorydatacenter_compute_allocation/
Directorydata/
- ai_labs.csv
- data_centers.csv
- diversity_caps.csv
- gpu_pools.csv
- lab_metrics.csv
- margin_floors.csv
- power_envelope_levels.csv
- workload_dependencies.csv
- workload_gpu_compatibility.csv
- workload_utilization_fallback.csv
- workload_utilization_test.csv
- workload_utilization_train.csv
- workload_utilization_val.csv
- workloads.csv
- datacenter_compute_allocation.py
- pyproject.toml
- README.md
- runbook.md
What this template is for
A datacenter operator — a neocloud, a colocation provider, or a hyperscaler running its own campuses — has energized GPU pools across several sites and must decide which AI lab’s workloads get which GPUs in which pool right now. Frontier labs want long-running training reservations on the largest contiguous pools at anchor-customer rates; inference customers want a steady slice of right-sized GPUs; research orgs want bursty best-effort capacity. They all compete for the same physical pools under the same power envelope, and the operator has to keep three tensions in balance at once: staying under each campus’s substation-bound power cap, holding gross margin after energy and depreciation, and serving anchor customers without letting any one customer’s share grow into a concentration risk. Getting the allocation wrong strands expensive capacity — depreciation accruing on GPUs that sit idle — which is the operator’s biggest economic exposure.
This template models that allocation as a scenario sweep across three axes (power envelope, margin floor, and anchor-diversity cap, for 24 combinations in all), tracing the trade-off frontiers the operator actually reasons about: how much revenue a tighter margin floor costs, and how much a tighter anchor cap costs. The strictest combinations come back infeasible, which is itself the answer — a signal that those constraints cannot all be met at once. A follow-on demand overlay then replays the chosen plan under softening-demand scenarios so the operator can see stranded-capacity exposure before committing. It picks up where the energy_grid_planning template leaves off: that upstream template decides which campuses get built and energized, and the bundled data_centers.csv here is a snapshot of its approved campus set.
The reasoning approach chains four reasoners on one shared ontology — a graph neural network (GNN) predicting per-workload utilization, declarative hardware-compatibility and priority rules, dependency-graph PageRank scoring how much downstream work each workload unblocks, and a scenario-parameterized mixed-integer program (MIP) that makes the assignment — with each stage writing signals the next stage reads.
Who this is for
- DC operators (neocloud, colo, hyperscaler infra teams) allocating GPU capacity across multi-tenant pools
- Capacity planners deciding when a (margin floor, diversity cap, envelope) point justifies acquiring more megawatts
- Operations researchers exploring multi-reasoner pipelines that compose a heterogeneous GNN, graph centrality, declarative rules, and a scenario-parameterized MIP
- Developers extending RelationalAI templates across decision lineages (capacity planning → allocation → scheduling)
What you’ll build
- A 6-lab × 5-DC compute allocation MIP, indexed by a 3-axis Scenario Concept sweep (2 envelopes × 3 margin floors × 4 diversity caps = 24 cells)
- A heterogeneous-graph GNN that classifies per-workload utilization probability (will this workload actually use its allocated capacity?) — the operator’s real forward-looking signal, propagated through lab→workload, dep-DAG, and cross-lab co_dated edges
- Reverse-PageRank scoring on a workload dependency DAG (4-5 deep chains rooted at frontier pretrains)
- Declarative hardware-compatibility rules + lex-emulating priority weights (P0=100, P1=10, P2=1)
- Two headline Pareto frontiers (margin × revenue and diversity × revenue) at the 100% envelope, with the 85% envelope drawn as an overlay sensitivity band
- Per-cell results queryable directly from the ontology — realized revenue, energy + depreciation cost, anchor share, per-workload-type counts, binding-constraint signal (OPTIMAL vs INFEASIBLE)
What’s included
datacenter_compute_allocation.py— main script: 4 reasoner stages + Pareto reporting + ontology persistencedata/data_centers.csv— 5 DCs (snapshot of the upstream $300M-approved set)data/gpu_pools.csv— 28 pools across 5 DCs (H100 / H200 / GB200 mix), with per-GPUhourly_depreciation_rate(1.52 H200, $2.66 GB200; capex/3-year amortization) data/ai_labs.csv— 6 labs: 3 frontier anchors (Anthropic, OpenAI, xAI Internal) + 2 applied (Cohere Inference, Together AI Multi-Lab) + 1 research (Stability Open)data/workloads.csv— 110 workloads: 15 P0 pretrains (frontier labs, 256-1024 GPUs each, GB200/H200/H100 mix), 30 P1 finetune (Together AI + Cohere), 50 P1 inference (Cohere + Together), 15 P2 eval (Stability)data/workload_gpu_compatibility.csv— per-workload GPU-type allowlistdata/workload_dependencies.csv— 138 edges (~130 blocks + ~10 informs), DAG depth 4-5 rooted at frontier pretrainsdata/lab_metrics.csv— 365 days × 6 labs = 2,190 rows with cross-lab co-movementdata/workload_utilization_train.csv,_val.csv,_test.csv— GNN temporal binary-classification splits. Train: 7 historical months × 110 workloads = 770(workload, observation_date, is_high_utilization)observations. Val: 1 month × 110 = 110. Test: the current month × 110 = 110 observations with no label (every workload gets a probability for the upcoming period).data/workload_utilization_fallback.csv— Stage 1 fallback (per-workload deterministic probability for the current period, used when--no-gnnis set)data/power_envelope_levels.csv— 2 scenario rows (0.85 / 1.00)data/margin_floors.csv— 3 scenario rows (unconstrained / 80% / 85% gross margin post-depreciation)data/diversity_caps.csv— 4 scenario rows (none / anchor_max_70pct / anchor_max_50pct CoreWeave-target / anchor_max_40pct_with_type_floor)runbook.md— analyst paste-test walkthrough mapping each stage to its RAI agent skill + prompt
Prerequisites
Access
- A Snowflake account with the RelationalAI native app installed.
- A Snowflake user with permissions on the RAI native app and on
EXP_DATABASE(the schema for GNN experiment artifacts). - A prescriptive engine for Stage 4. Optionally a GPU-backed predictive engine for Stage 1.
Tools
- Python ≥ 3.10.
- RelationalAI Python SDK with the predictive submodule (
relationalai[gnn], providingrelationalai.semantics.reasoners.predictive) for Stage 1’s GNN. Without it, Stage 1 falls back to the precomputedworkload_utilization_fallback.csvautomatically.
One-time Snowflake setup for GNN experiment artifacts
The predictive reasoner writes training artifacts (model checkpoints, metrics, predictions) to a Snowflake schema that the RELATIONALAI native app must own. Pick a database you control, then create + grant the schema:
CREATE DATABASE IF NOT EXISTS DATACENTER_ENRICHMENT;CREATE SCHEMA IF NOT EXISTS DATACENTER_ENRICHMENT.EXPERIMENTS;GRANT USAGE ON DATABASE DATACENTER_ENRICHMENT TO APPLICATION RELATIONALAI;GRANT USAGE ON SCHEMA DATACENTER_ENRICHMENT.EXPERIMENTS TO APPLICATION RELATIONALAI;GRANT CREATE EXPERIMENT ON SCHEMA DATACENTER_ENRICHMENT.EXPERIMENTS TO APPLICATION RELATIONALAI;GRANT CREATE MODEL ON SCHEMA DATACENTER_ENRICHMENT.EXPERIMENTS TO APPLICATION RELATIONALAI;Set EXP_DATABASE at the top of datacenter_compute_allocation.py to that database (default: DATACENTER_ENRICHMENT).
Quickstart
-
Download the template and extract it:
Terminal window curl -O https://docs.relational.ai/templates/zips/v1/datacenter_compute_allocation.zipunzip datacenter_compute_allocation.zipcd datacenter_compute_allocation -
Create a virtual environment and activate it:
Terminal window python -m venv .venvsource .venv/bin/activatepython -m pip install --upgrade pip -
Install dependencies:
Terminal window python -m pip install . -
Configure your RAI connection:
Terminal window rai init -
Run the template:
Terminal window python datacenter_compute_allocation.py -
Skip the GNN and load
workload_utilization_fallback.csvdirectly (useful for fast iteration on Stages 2–4):Terminal window python datacenter_compute_allocation.py --no-gnn -
Expected output (abbreviated):
STAGE 4: PRESCRIPTIVE -- compute allocation MIP (24-cell sweep)Termination status: OPTIMALPer-cell summary (24 cells: 16 optimal, 8 infeasible)AllocationPlan singleton (100pct / unconstrained / none):110 workloads assigned, $25.28M revenue, 83% margin, 95% anchor shareFull stage-by-stage log (GNN distribution, PageRank scores, both Pareto frontiers, the demand-scenario overlay, and expected per-cell behavior) is in
runbook.md. GNN probabilities depend on seed and predictive engine; the--no-gnnfallback values are deterministic. Stage 4 reachesOPTIMALin seconds with Gurobi; the open-source HiGHS default returns a feasible solution at thetime_limit_sec=240wall and may surfaceTIME_LIMIT(signal, not failure, perrai-prescriptive-results). SetSOLVERindatacenter_compute_allocation.pyto switch.
Template structure
datacenter_compute_allocation/ datacenter_compute_allocation.py # Main script (4 chained reasoning stages + ontology persistence) data/ data_centers.csv # 5 DCs (snapshot of upstream $300M-approved set) gpu_pools.csv # 28 GPU pools across 5 DCs (H100/H200/GB200 mix) ai_labs.csv # 6 labs (3 frontier, 2 applied, 1 research) workloads.csv # 110 workloads (15 P0 / 80 P1 / 15 P2) workload_gpu_compatibility.csv # Per-workload GPU-type allowlist workload_dependencies.csv # 138 directed edges in the dependency DAG lab_metrics.csv # 2,190 = 365 days × 6 labs daily KPIs workload_utilization_train.csv # GNN train split: 7 historical months × 110 = 770 obs workload_utilization_val.csv # GNN val split: 1 month × 110 = 110 obs workload_utilization_test.csv # GNN test: current period × 110, no label workload_utilization_fallback.csv # Stage 1 fallback (GNN-shape per-workload p) power_envelope_levels.csv # 2 scenario rows (envelope multiplier) margin_floors.csv # 3 scenario rows (gross margin floors) diversity_caps.csv # 4 scenario rows (anchor concentration caps) README.md # this file runbook.md # multi-reasoner agent-skill walkthrough pyproject.toml # dependenciesStart here: run python datacenter_compute_allocation.py for the full four-stage chain end to end (add --no-gnn to skip Stage 1’s GNN and load the deterministic fallback), or follow runbook.md to rebuild it step by step.
Sample data
The bundled data is synthetic and illustrative — designed to teach the four-stage reasoning flow on a Snowflake-connected RAI account, not to match a specific operator’s fleet. The five campuses, six labs, GPU generations, and workload mix are shaped so every stage does visible work and the Pareto frontiers and infeasible cells show up as intended.
data_centers.csv(5 rows) — campuses (snapshot of the upstream $300M-approved set), each with an approved power cap, PUE, and energy cost.gpu_pools.csv(28 rows) — GPU pools across the 5 campuses (H100 / H200 / GB200 mix), each with a per-GPU hourly depreciation rate.ai_labs.csv(6 rows) — labs: 3 frontier anchors, 2 applied, 1 research, each with a contract tier and anchor flag.workloads.csv(110 rows) — training, fine-tune, inference, and eval workloads with GPU and memory requirements and a strategic dollar value.workload_gpu_compatibility.csv— the per-workload GPU-type allowlist (multi-valued, so it lives as its own concept).workload_dependencies.csv(138 rows) — directed dependency edges (blocks/informs), DAG depth 4-5 rooted at frontier pretrains.lab_metrics.csv(2,190 rows) — 365 days x 6 labs of daily lab-activity KPIs with cross-lab co-movement, the GNN’s temporal feature source.workload_utilization_train.csv/_val.csv/_test.csv— the GNN’s temporal binary-classification splits (7 historical months / 1 validation month / current period, no label).workload_utilization_fallback.csv— the deterministic Stage 1 fallback used when--no-gnnis set.power_envelope_levels.csv/margin_floors.csv/diversity_caps.csv— the three scenario-axis tables (2 / 3 / 4 rows) that index the 24-cell sweep.
Model overview
One shared ontology threads all four stages: each stage reads concepts and properties earlier stages wrote and writes new ones for downstream stages.
- Key entities:
DataCenterRequest(a campus),GpuPool,AILab,Workload,WorkloadGpuCompatandWorkloadDependency(workload relationships as their own concepts), the three scenario conceptsPowerEnvelopeLevel/MarginFloor/DiversityCap, plus the derivedCompatibilityandAssignment(decision space), theAllocationPlansingleton, and theDemandScenariooverlay. - Primary identifiers: string ids on
DataCenterRequest,GpuPool,AILab, andWorkload; the scenario concepts are identified byname;WorkloadGpuCompat,WorkloadDependency,Compatibility, andAssignmentcarry composite keys. - Important invariants:
envelope_multiplier,fraction, andanchor_max_shareare fractions; a negative sentinel onMarginFloor.fraction/DiversityCap.anchor_max_sharemarks the unconstrained row (its constraint is skipped); theis_high_utilizationlabel is binary;Assignment.x_assigndecision variables are binary; power / cost / GPU counts are non-negative.
For the full concept and property definitions — the entity concepts, the three scenario concepts, and the derived Compatibility / Assignment / AllocationPlan / DemandScenario concepts each stage writes — see datacenter_compute_allocation.py; runbook.md builds them step by step with the RAI skills.
How it works
Stage 1: Predictive — per-workload utilization-probability GNN
Binary node classification on Workload: predict the probability that each workload will be high-utilization (actually use its allocated capacity at high duty cycle) vs stall or be repaced. Stranded capacity (depreciation accruing without offsetting revenue) is the operator’s biggest economic exposure, so this is the chain’s load-bearing forward-looking signal.
The graph is heterogeneous with three cross-concept edge types: a LabMetric → Workload edge that carries lab-side recent activity into a workload’s prediction neighborhood (a workload owned by a fast-ramping lab inherits its signal); the WorkloadDependency.blocks DAG shared with Stage 3, so a workload downstream of a high-utilization gating pretrain inherits signal through the dependency chain; and a load-bearing LabMetric → LabMetric edge between cross-lab same-date pairs sharing a workload type, which carries cross-lab co-movement (funding waves, supply shocks).
The task tables carry per-(workload, observation_date) historical labels. Each row is one monthly utilization observation, so the same workload appears in 7 training rows (one per historical month), giving the GNN per-period variety to learn from instead of a single label per entity — Train (7 months × 110 = 770) / Val (1 month × 110 = 110) / Test (current month × 110, no label, so every workload gets a probability for the upcoming period). The task relationships use at time slots and the GNN runs with has_time_column=True so LabMetric features are time-aligned with each (workload, month) prediction. The positive-class probability is bound back to the ontology as Workload.utilization_probability and consumed by Stage 4’s objective.
Stage 2: Rules — eligibility + priority tiers + Compatibility precompute
Two rule families run on the workload-pool product. The eligibility rule joins through the WorkloadGpuCompat(workload, gpu_type) composite-key concept (per the GNN-node FD trap: multi-valued allowlists on a GNN-node concept must live as their own concept, not as a Workload Relationship) and enforces both the memory check and the GPU-type allowlist. Eligible pairs are then materialized into a Compatibility(workload, gpu_pool) concept — a precompute that keeps the Stage 4 MIP linear and small. Priority tier classification reads AILab.contract_tier and writes both Workload.priority_tier (string) and Workload.priority_weight (100 / 10 / 1) — the numeric form is what Stage 4’s objective consumes.
Stage 3: Graph — reverse-PageRank for downstream gating
A directed Workload graph is built with edge reversal (successor → predecessor) so a node’s PageRank accumulates flow from everything it gates downstream. The resulting score lands on Workload.gating_score. Workloads absent from the DAG get a backstop gating_score = 1.0 so they enter the Stage 4 objective product without zeroing out.
Stage 4: Prescriptive — assignment MIP under a 3D Scenario sweep
The decision variable Assignment.x_assign is a binary variable indexed by the three scenario concepts (PowerEnvelopeLevel, MarginFloor, DiversityCap), so the MIP solves all 24 cells in a single pass. Constraint families run one per scenario axis:
- C1+C2 at-most-one per workload per cell (soft P0; priority_weight = 100 drives P0 from the objective).
- C3 per-pool GPU-count capacity.
- C4 per-DC power envelope:
Σ gpu_count × power_per_gpu × pue ≤ approved_mw × 1000 × envelope_multiplier. - C5 linearized gross-margin floor (skipped when
MarginFloor.fraction < 0for the unconstrained row). - C6 anchor-concentration cap (skipped when
DiversityCap.anchor_max_share < 0). - C7 workload-type floor (skipped unless
DiversityCap.workload_type_floor >= 0).
The objective is what makes this a real chain rather than four disjoint reasoners: it maximizes the sum over chosen assignments of priority_weight (Stage 2) × gating_score (Stage 3) × utilization_probability (Stage 1) × strategic_value_usd (raw $ baseline) — every prior stage’s signal enters the same product. After the solve, the per-cell summary is assembled with ontology-side aggregation (no pandas groupby). The chosen baseline cell (100pct / unconstrained / none) is persisted as a singleton AllocationPlan concept with revenue / total_cost / margin / anchor_share / n_assigned / status / binding_axis, plus an Assignment.is_chosen flag over its decision rows.
Finally, a DemandScenario overlay replays the locked-in plan under four risk scenarios. Anchor (P0) revenue is treated as contractual (factor 1.0) — the operator gets paid for the seat regardless of utilization — while opportunistic (P1/P2) seats realize only the scenario factor of their assigned revenue. The overlay reports realized and stranded revenue per scenario and persists each as a DemandScenarioOutlook(scenario) concept row, so the stranded-capacity exposure of the chosen plan is queryable as ontology, not just printed.
See datacenter_compute_allocation.py for the implementation and runbook.md to reproduce each stage step by step with the RAI skills.
Customize this template
Use your own data
- Add or remove campuses by editing
data_centers.csv. The bundled set is a snapshot of the upstream $300M-approved campuses; re-runningenergy_grid_planningat a differentInvestmentLevelproduces a different approved set you can copy in. - Adjust the lab roster in
ai_labs.csv. The anchor flag drives theDiversityCapconstraint. - Replace the GNN with the deterministic fallback by setting
--no-gnnand editingworkload_utilization_fallback.csvdirectly.
Tune parameters
- Tune the scenario axes: edit
power_envelope_levels.csv,margin_floors.csv, anddiversity_caps.csv. Adding rows scales the cell count multiplicatively. - Change the priority spread by editing the weights in the
# Stage 2: Rulessection (default 100 / 10 / 1).
Extend the model
- Add a fourth GNN edge type, a new dependency relation, or extra scenario axes; the accretive-ontology pattern means downstream stages pick up new signals without restructuring.
Scale up / productionize
- Tune the Stage 4 solve by adjusting the solver name or
time_limit_secinproblem.solve(...). The default works on any prescriptive engine; faster commercial solvers (when licensed) can reach OPTIMAL in a fraction of the time. - Replace the bundled CSVs with CDC ingestion from the operator’s inventory, telemetry, and contract systems; the ontology shape is independent of the load pipeline.
Troubleshooting
ModuleNotFoundError: relationalai.semantics.reasoners.predictive
The installed relationalai predates the predictive submodule. Pin a release that ships the GNN, or run with --no-gnn to load the deterministic fallback CSV.
Database does not exist or the GNN RelationalAI Native App lacks permissions
The DATACENTER_ENRICHMENT.EXPERIMENTS schema is not provisioned with the required grants. Run the four GRANT statements in Prerequisites.
TransactionAbortedException during gnn.fit()
Runtime error in one of the GNN’s exported tables — usually a multi-valued Relationship on a GNN-node concept tripping the engine’s functional-dependency check. Pull the engine-side root cause:
snow sql -q "CALL RELATIONALAI.API.GET_TRANSACTION_ARTIFACTS('<txn_id>', OBJECT_CONSTRUCT())"Then download the problems.json artifact. The most common fix is moving the offending multi-valued relationship to its own composite-id Concept (mirrors how Compatibility(workload, gpu_pool) and WorkloadGpuCompat(workload, gpu_type) are themselves Concepts, not Workload Relationships).
Solver returns INFEASIBLE for the global solve
Some cell’s constraints are mutually unsatisfiable AND the C1 P0 commitment is a hard ==1. C1 should be soft <=1; the priority_weight=100 in the objective drives P0 to 1 wherever feasible. Per-cell infeasibility then surfaces as INFEASIBLE in the per-cell summary, not as global failure.
Solver returns TIME_LIMIT with sensible per-cell results
Expected at the default time_limit_sec=240. Per-cell numbers remain valid with TIME_LIMIT (rai-prescriptive-results: TIME_LIMIT is signal, not error). Raise time_limit_sec or switch to a faster solver in the # Stage 4: Prescriptive section if you need a tighter MIP gap.
Per-cell summary shows all cells INFEASIBLE
Likely an upstream constraint (C5 / C6 / C7) firing at unintended cells. Check that the -1.0 NULL-handling sentinels are filtered correctly in C5 / C6 / C7’s model.where(... .fraction >= 0.0 ...) clauses.
Per-workload utilization probabilities do not match the README's distribution
GNN trained with non-default seed/epochs, or --no-gnn fallback used. The fallback values in workload_utilization_fallback.csv are the deterministic reference (the latent probability used to generate the synthetic labels); these are what downstream stages consume when --no-gnn is set.
What this template abstracts
This is a snapshot allocator, not a production scheduler. Workloads are assigned to a single (Workload, GpuPool) slot — gang scheduling, topology-aware placement, time-block reservations, and elastic sizing are out of scope. Utilization is modeled as a per-workload point estimate; per-workload (p10, p50, p90) bands would be a natural extension, with the scenario sweep providing the aggregate-level approximation today.
Learn more
- Upstream template: energy_grid_planning — the multi-year capex / interconnect decision that approves which campuses get built and energized.
- Peer multi-reasoner: telco_network_recovery — same accretive-chain pattern in a tower-upgrade context, using a binary-classification GNN + three-branch rule + PageRank + MIP.
- RelationalAI documentation: https://docs.relational.ai
Support
- File issues at https://github.com/RelationalAI/templates/issues.
- Reach the RAI team via the support channels listed at https://relational.ai.