SkyDiscover: A Flexible Framework for AI-Driven Scientific and Algorithmic Discovery

# SPDX-License-Identifier: Apache-2.0
"""
Expert parallelism load balancer (EPLB) for vLLM.

This module implements the core rearrangement algorithm.

The rearrangement algorithm is adapted from
[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).

Please find at [#12](https://github.com/deepseek-ai/EPLB/issues/12) an example
on how the EPLB algorithm works.
"""

# EVOLVE-BLOCK-START

import torch

def _balanced_round_robin_pack(costs: torch.Tensor, num_packs: int, items_per_pack: int) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Equal-cardinality balanced packing.
    Approach:
      - Sort per-instance costs descending.
      - In round r, take next num_packs items and assign to packs sorted by current load ascending (stable).
    Returns:
      pack_id [M], rank_in_pack [M] where M = num_packs * items_per_pack.
    """
    M = costs.numel()
    assert M == num_packs * items_per_pack
    order = torch.argsort(costs, descending=True)
    pack_id = torch.empty(M, dtype=torch.int64)
    rank_in_pack = torch.empty(M, dtype=torch.int64)
    loads = torch.zeros(num_packs, dtype=costs.dtype)

    for r in range(items_per_pack):
        start, end = r * num_packs, (r + 1) * num_packs
        idx = order[start:end]
        gpu_order = torch.argsort(loads, stable=True)
        pack_id[idx] = gpu_order
        rank_in_pack[idx] = r
        loads.index_add_(0, gpu_order, costs[idx])
    return pack_id, rank_in_pack

def _apportion_divisor(tokens: torch.Tensor, total_replicas: int) -> torch.Tensor:
    """
    Divisor water-filling apportionment (Webster-like).
    - Guarantee at least one replica per expert.
    - Binary search threshold T so sum_i max(0, ceil(w_i/T)-1) matches extras.
    - Resolve rounding with largest fractional remainders; final exact-sum fixup.
    Complexity per layer: O(E log U).
    """
    E = tokens.numel()
    counts = torch.ones(E, dtype=torch.int64)
    extras = int(total_replicas - E)
    if extras <= 0:
        return counts

    w = tokens.to(torch.float64).clamp_min(0.0).cpu()
    if float(w.max().item()) == 0.0:
        # Uniformly distribute extras on all-zero weights
        base, rem = divmod(extras, E)
        counts += base
        if rem:
            counts[:rem] += 1
        return counts

    lo, hi = 0.0, float(w.max().item()) + 1e-12
    for _ in range(30):
        T = (lo + hi) * 0.5
        s = int(torch.clamp(torch.ceil(w / T) - 1.0, min=0.0).sum().item())
        if s > extras:
            lo = T
        else:
            hi = T

    x = torch.clamp(w / hi - 1.0, min=0.0)
    extras_int = torch.floor(x).to(torch.int64)
    counts = extras_int + 1
    deficit = extras - int(extras_int.sum().item())

    if deficit > 0:
        frac = (x - torch.floor(x))
        bias = (w / (w.sum() + 1e-12)) * 1e-12  # stable tie-break toward heavier experts
        frac = (frac + bias).to(torch.float64)
        k = min(deficit, E)
        if k:
            idx = torch.topk(frac, k=k, sorted=False).indices
            counts.index_add_(0, idx, torch.ones_like(idx, dtype=torch.int64))
    elif deficit < 0:
        removable = counts > 1
        if int(removable.sum().item()) > 0:
            vals = torch.where(removable, x, torch.full_like(x, float("inf")))
            k = min(-deficit, int(removable.sum().item()))
            if k:
                idx = torch.topk(-vals, k=k, sorted=False).indices
                counts.index_add_(0, idx, -torch.ones_like(idx, dtype=torch.int64))

    diff = int(total_replicas - int(counts.sum().item()))
    if diff != 0:
        order = torch.argsort(w, descending=(diff > 0))
        step = 1 if diff > 0 else -1
        i = 0
        while diff != 0 and i < E:
            j = int(order[i].item())
            if step > 0 or counts[j] > 1:
                counts[j] += step
                diff -= step
            i += 1
    return counts.to(torch.int64)

def rebalance_experts(
    weight: torch.Tensor,
    num_replicas: int,
    num_groups: int,
    num_nodes: int,
    num_gpus: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Global fast rebalance:
      1) Per-layer replica apportionment via divisor water-filling (>=1 per expert).
      2) Place per-replica instances onto GPUs by equal-cardinality balanced round-robin.
    Returns:
      physical_to_logical_map [L,R], logical_to_physical_map [L,E,X], expert_count [L,E].
    """
    L, E = weight.shape
    assert num_replicas % num_gpus == 0
    assert num_replicas >= E
    per_gpu = num_replicas // num_gpus

    w = weight.float().cpu().contiguous()
    device = torch.device("cpu")

    phy2log = torch.empty(L, num_replicas, dtype=torch.int64, device=device)
    phyrank = torch.empty(L, num_replicas, dtype=torch.int64, device=device)
    logcnt = torch.zeros(L, E, dtype=torch.int64, device=device)

    for layer in range(L):
        wl = w[layer]
        counts = _apportion_divisor(wl, num_replicas)
        logcnt[layer] = counts

        arange_log = torch.arange(E, dtype=torch.int64)
        inst2log = torch.repeat_interleave(arange_log, counts)
        starts = torch.cumsum(counts, 0) - counts
        inst_rank = torch.arange(num_replicas, dtype=torch.int64) - torch.repeat_interleave(starts, counts)

        inst_cost = (wl[inst2log] / counts[inst2log].to(wl.dtype)).contiguous()
        pack_id, rank = _balanced_round_robin_pack(inst_cost, num_packs=num_gpus, items_per_pack=per_gpu)
        slot = pack_id * per_gpu + rank

        phy2log[layer, slot] = inst2log
        phyrank[layer, slot] = inst_rank

    redundant = num_replicas - E
    maxlogcnt = redundant + 1
    log2phy = torch.full((L, E, maxlogcnt), -1, dtype=torch.int64, device=device)
    scatter_idx = phy2log * maxlogcnt + phyrank
    log2phy.view(L, -1).scatter_(
        -1,
        scatter_idx,
        torch.arange(num_replicas, dtype=torch.int64, device=device).expand(L, -1),
    )
    return phy2log, log2phy, logcnt

# EVOLVE-BLOCK-END

__all__ = ["rebalance_experts"]

GPU_MEM_SIZE = 80 # GB

# EVOLVE-BLOCK-START

def compute_model_placement(gpu_num, models):
    """Minimize max KVPR via binary search on T. For target T, give each model weight w=a+T*s (a=req_rate/slo, s=size) and each GPU capacity T*M; pack with Best-Fit-Decreasing and return placement for minimal feasible T."""
    M=GPU_MEM_SIZE; EPS=1e-12
    it=[(m, m.req_rate/max(m.slo,EPS), m.model_size) for m in models]
    if any(s>M for _,_,s in it) or sum(s for _,_,s in it)>gpu_num*M:
        raise ValueError("Unable to place model")
    def pack(T,ret=False):
        cap=[T*M]*gpu_num; P=[[] for _ in range(gpu_num)]
        ws=sorted(((m,a+T*s) for m,a,s in it), key=lambda x:x[1], reverse=True)
        for m,w in ws:
            bi=-1; br=float('inf')
            for i in range(gpu_num):
                r=cap[i]-w
                if r>=-1e-12 and r<br:
                    bi=i; br=r
            if bi<0: return None if ret else False
            cap[bi]-=w; P[bi].append(m)
        return {i:P[i]} if ret else True
    hi=1.0; tries=0
    while not pack(hi):
        hi*=2.0; tries+=1
        if tries>60: raise ValueError("Unable to place model")
    lo=0.0
    for _ in range(35):
        mid=(lo+hi)/2.0
        if pack(mid): hi=mid
        else: lo=mid
    res=pack(hi,True)
    if res is None: raise ValueError("Unable to place model")
    return res

# EVOLVE-BLOCK-END

if __name__ == "__main__":
    # Test the algorithm

    from evaluator import generate_test_gpu_models
    from evaluator import calculate_kvcache_pressure
    from evaluator import safe_float
    import numpy as np

    test_cases = generate_test_gpu_models()
    all_kvpr = []
    for i, (gpu_num, gpu_models) in enumerate(test_cases):

        results = compute_model_placement(gpu_num, gpu_models)
        max_kvpr = calculate_kvcache_pressure(results)
        all_kvpr.append(safe_float(max_kvpr))

    avg_kvpr = np.mean(all_kvpr)
    if avg_kvpr != 0:
        avg_kvpr = 1.0 / avg_kvpr

    print(f"Max KVPR: {avg_kvpr:.3f}")

# EVOLVE-BLOCK-START
import networkx as nx
import json
import os
import pandas as pd
from typing import Dict, List

def search_algorithm(src, dsts, G, num_partitions):
    """
    Ensemble Steiner Tree Search with Frequency-Based Candidate Selection.
    1. Preprocessing: Cleans graph, computes APSP, and identifies high-frequency intermediate nodes
       appearing in shortest paths between terminals.
    2. Multi-Start Search: Generates a population of trees using:
       - Deterministic Takahashi-Matsuyama.
       - Randomized Shortest Path Trees (perturbed weights).
    3. Refinement: Applies Iterative MSA to the best candidate sets to optimize the topology.
    4. Selection: Returns the lowest-cost tree found.
    """
    import random
    from collections import Counter

    # 1. Graph Cleanup & Precomputation
    h = G.copy()
    valid_edges = []
    for u, v, d in h.edges(data=True):
        if d.get("cost") is not None and d.get("throughput", 0) > 0:
            valid_edges.append((u, v, d))

    h_clean = nx.DiGraph()
    h_clean.add_edges_from(valid_edges)
    h_clean.add_nodes_from(G.nodes())

    try:
        apsp = dict(nx.all_pairs_dijkstra(h_clean, weight="cost"))
    except:
        return BroadCastTopology(src, dsts, num_partitions)

    if src not in apsp:
        return BroadCastTopology(src, dsts, num_partitions)

    terminals = set(dsts) | {src}

    # 2. Candidate Identification (Frequency Analysis)
    node_freq = Counter()
    terminal_list = list(terminals)
    pairs = []
    if len(terminal_list) > 1:
        for _ in range(100):
            u, v = random.sample(terminal_list, 2)
            if u != v: pairs.append((u, v))

    for u, v in pairs:
        if u in apsp and v in apsp[u][1]:
            path = apsp[u][1][v]
            node_freq.update(path)

    freq_candidates = [n for n, c in node_freq.most_common(20) if n not in terminals]

    best_tree = None
    min_cost = float('inf')

    def get_tree_cost(t):
        return sum(d.get("cost", 0) for _, _, d in t.edges(data=True))

    def refine_candidates(initial_cands):
        curr_cands = set(initial_cands)
        local_best = None
        local_min = float('inf')

        for _ in range(3):
            active = list(terminals | curr_cands)
            metric_g = nx.DiGraph()

            for u in active:
                if u not in apsp: continue
                targets = apsp[u][0]
                for v in active:
                    if u != v and v in targets:
                        metric_g.add_edge(u, v, weight=targets[v])

            if src in metric_g:
                metric_g.remove_edges_from(list(metric_g.in_edges(src)))

            if not metric_g.nodes(): break

            try:
                msa = nx.minimum_spanning_arborescence(metric_g, attr="weight")
            except: break

            leaves = [n for n in msa.nodes() if msa.out_degree(n) == 0 and n not in terminals]
            msa.remove_nodes_from(leaves)

            phys = nx.DiGraph()
            possible = True
            for u, v in msa.edges():
                if u in apsp and v in apsp[u][1]:
                    path = apsp[u][1][v]
                    nx.add_path(phys, path)
                    for x, y in zip(path, path[1:]):
                        if h_clean.has_edge(x, y):
                            phys.add_edge(x, y, **h_clean[x][y])
                        else: possible = False
                else: possible = False

            if not possible: break

            c = get_tree_cost(phys)
            if c < local_min:
                local_min = c
                local_best = phys

            branching = {n for n in phys.nodes() if phys.out_degree(n) > 1}
            new_cands = (branching | set(phys.nodes())) - terminals
            if new_cands.issubset(curr_cands): break
            curr_cands = new_cands

        return local_best, local_min

    # Strategy A: Takahashi-Matsuyama
    tm_cands = set()
    tm_reached = {src}
    tm_unreached = set(dsts) - {src}
    while tm_unreached:
        best_n, best_p, min_d = None, None, float('inf')
        for u in tm_reached:
            if u not in apsp: continue
            for v in tm_unreached:
                if v in apsp[u][0] and apsp[u][0][v] < min_d:
                    min_d = apsp[u][0][v]
                    best_n, best_p = v, u
        if best_n:
            path = apsp[best_p][1][best_n]
            tm_cands.update(path)
            tm_reached.update(path)
            tm_unreached.remove(best_n)
        else: break

    t, c = refine_candidates(tm_cands - terminals)
    if t and c < min_cost:
        min_cost, best_tree = c, t

    # Strategy B: Randomized Perturbation
    for i in range(8):
        perturb = 0.05 + (i * 0.05)
        p_G = nx.DiGraph()
        for u, v, d in h_clean.edges(data=True):
            w = d.get("cost", 0) * random.uniform(1.0 - perturb, 1.0 + perturb)
            p_G.add_edge(u, v, weight=w)

        trial_cands = set(freq_candidates)
        try:
            p_dists, p_paths = nx.single_source_dijkstra(p_G, src, weight="weight")
            for d in dsts:
                if d in p_paths:
                    trial_cands.update(p_paths[d])
        except: pass

        t, c = refine_candidates(trial_cands - terminals)
        if t and c < min_cost:
            min_cost, best_tree = c, t

    bc = BroadCastTopology(src, dsts, num_partitions)
    final = best_tree if best_tree else h_clean

    for dst in dsts:
        if dst == src: continue
        edges = []
        if final.has_node(dst) and nx.has_path(final, src, dst):
            try:
                p = nx.shortest_path(final, src, dst, weight="cost")
                edges = [[u, v, final[u][v]] for u, v in zip(p, p[1:])]
            except: pass

        if not edges and src in apsp and dst in apsp[src][1]:
            p = apsp[src][1][dst]
            edges = [[u, v, h_clean[u][v]] for u, v in zip(p, p[1:])]

        if edges:
            for i in range(num_partitions):
                bc.set_dst_partition_paths(dst, i, edges)

    return bc

class SingleDstPath(Dict):
    partition: int
    edges: List[List]  # [[src, dst, edge data]]

class BroadCastTopology:
    def __init__(self, src: str, dsts: List[str], num_partitions: int = 4, paths: Dict[str, SingleDstPath] = None):
        self.src = src  # single str
        self.dsts = dsts  # list of strs
        self.num_partitions = num_partitions

        # dict(dst) --> dict(partition) --> list(nx.edges)
        # example: {dst1: {partition1: [src->node1, node1->dst1], partition 2: [src->dst1]}}
        if paths is not None:
            self.paths = paths
            self.set_graph()
        else:
            self.paths = {dst: {str(i): None for i in range(num_partitions)} for dst in dsts}

    def get_paths(self):
        print(f"now the set path is: {self.paths}")
        return self.paths

    def set_num_partitions(self, num_partitions: int):
        self.num_partitions = num_partitions

    def set_dst_partition_paths(self, dst: str, partition: int, paths: List[List]):
        """
        Set paths for partition = partition to reach dst
        """
        partition = str(partition)
        self.paths[dst][partition] = paths

    def append_dst_partition_path(self, dst: str, partition: int, path: List):
        """
        Append path for partition = partition to reach dst
        """
        partition = str(partition)
        if self.paths[dst][partition] is None:
            self.paths[dst][partition] = []
        self.paths[dst][partition].append(path)

def make_nx_graph(cost_path=None, throughput_path=None, num_vms=1):
    """
    Default graph with capacity constraints and cost info
    nodes: regions, edges: links
    per edge:
        throughput: max tput achievable (gbps)
        cost: $/GB
        flow: actual flow (gbps), must be < throughput, default = 0
    """
    # Use relative path from this file's location
    current_dir = os.path.dirname(os.path.abspath(__file__))

    if cost_path is None:
        cost = pd.read_csv(os.path.join(current_dir, "profiles/cost.csv"))
    else:
        cost = pd.read_csv(cost_path)

    if throughput_path is None:
        throughput = pd.read_csv(os.path.join(current_dir, "profiles/throughput.csv"))
    else:
        throughput = pd.read_csv(throughput_path)

    G = nx.DiGraph()
    for _, row in throughput.iterrows():
        if row["src_region"] == row["dst_region"]:
            continue
        G.add_edge(row["src_region"], row["dst_region"], cost=None, throughput=num_vms * row["throughput_sent"] / 1e9)

    for _, row in cost.iterrows():
        if row["src"] in G and row["dest"] in G[row["src"]]:
            G[row["src"]][row["dest"]]["cost"] = row["cost"]

    # some pairs not in the cost grid
    no_cost_pairs = []
    for edge in G.edges.data():
        src, dst = edge[0], edge[1]
        if edge[-1]["cost"] is None:
            no_cost_pairs.append((src, dst))
    print("Unable to get costs for: ", no_cost_pairs)

    return G

# EVOLVE-BLOCK-END

# Helper functions that won't be evolved
def create_broadcast_topology(src: str, dsts: List[str], num_partitions: int = 4):
    """Create a broadcast topology instance"""
    return BroadCastTopology(src, dsts, num_partitions)

def run_search_algorithm(src: str, dsts: List[str], G, num_partitions: int):
    """Run the search algorithm and return the topology"""
    return search_algorithm(src, dsts, G, num_partitions)