Repository: ray-project/ray
File: ci/ray_ci/bazel_sharding.py (Python)

Code before:

def allocate_slots_to_shards(
    rules_grouped_by_time: List[Tuple[float, List[BazelRule]]],
    count: int,
) -> Tuple[List[dict], float]:
    """
    Allocate test slots to shards using least-loaded strategy.

    This only determines how many tests of each size go to each shard,
    without assigning specific tests. This preserves load balancing while
    allowing tests to be assigned in name order later.

    Returns:
        A tuple of (shard_slots, optimum) where shard_slots[i] is a dict
        mapping timeout -> number of slots for shard i.
    """
    shard_times = [0.0] * count
    shard_slots = [defaultdict(int) for _ in range(count)]

    # The theoretical optimum we are aiming for. Note that this may be unattainable
    # as it doesn't take into account that tests are discrete and cannot be split.
    # This is however fine, because it should only serve as a guide which shard to
    # add the next test to.
    optimum = (
        sum(timeout * len(rules) for timeout, rules in rules_grouped_by_time) / count
    )

    for timeout, rules in rules_grouped_by_time:
        for _ in range(len(rules)):
            # Find the least-loaded shard below optimum, or closest to optimum
            best_below_idx = None
            best_below_time = None
            best_above_idx = None
            best_above_time = None

            for i, shard_time in enumerate(shard_times):
                time_with_slot = shard_time + timeout
                # Among shards below optimum, pick the least loaded
                if time_with_slot < optimum:
                    if best_below_idx is None or shard_time < best_below_time:
                        best_below_idx = i
                        best_below_time = shard_time
                # Among shards at/above optimum, pick closest to optimum
                elif best_above_idx is None or time_with_slot < best_above_time:
                    best_above_idx = i
                    best_above_time = time_with_slot

            best_idx = best_below_idx if best_below_idx is not None else best_above_idx
            shard_slots[best_idx][timeout] += 1
            shard_times[best_idx] += timeout


Review comment (suggestion):
we do not need this right? you are just always putting tests into the shard that has the min sum?

can we just remove the optimum concept now? this algorithm does not have anything to do with the expected optimum.

Code after:

def allocate_slots_to_shards(
    rules_grouped_by_time: List[Tuple[int, List[BazelRule]]],
    count: int,
) -> List[Dict[int, int]]:
    """
    Allocate test slots to shards using least-loaded strategy.

    This only determines how many tests of each size go to each shard,
    without assigning specific tests. This preserves load balancing while
    allowing tests to be assigned in name order later.
    """
    shard_times = [0] * count
    shard_slots = [defaultdict(int) for _ in range(count)]

    for timeout, rules in rules_grouped_by_time:
        for _ in range(len(rules)):
            # Always pick the least-loaded shard
            best_idx = min(range(count), key=lambda i: shard_times[i])
            shard_slots[best_idx][timeout] += 1
            shard_times[best_idx] += timeout

    return shard_slots


def get_rules_for_shard_naive(
    rules_grouped_by_time: List[Tuple[int, List[BazelRule]]], index: int, count: int
) -> List[str]:
    """Create shards by assigning the same number of rules to each shard."""
    all_rules = []
    for _, rules in rules_grouped_by_time:
        all_rules.extend(rules)
    shard = sorted(all_rules)[index::count]
    return [rule.name for rule in shard]


def get_rules_for_shard_optimal(
    rules_grouped_by_time: List[Tuple[int, List[BazelRule]]], index: int, count: int
) -> List[str]:
    """Creates shards by trying to make sure each shard takes around the same time.

    Uses a two-phase approach:
    1. Allocate slots to shards using least-loaded strategy (determines capacity)
    2. Fill slots with tests in name order (preserves test clustering)

    This ensures no empty shards while keeping tests of the same size contiguous,
    making it easier to locate specific tests.

    ``rules_grouped_by_time`` is expected to be a list of tuples of
    (timeout in seconds, list of rules) sorted by timeout descending.
    """