The name of mslpr is kind of confused here. Just use the full name is fine.

The readability of this function could be improved, example code refactored with GPT:

def dcp_make_buckets_equal(
    sample_seqlens: List[Tuple[int, int]],
    compute_estimator: Callable,
    max_seq_len_per_rank: int,
    min_cp_size: int = 1,
	bucket_overfill_factor = 1.1,
) -> List[deque]:
    """Split samples into buckets of roughly equal work, one per unique CP size
    Args:
        sample_seqlens (List[Tuple[int, int]]): (sample_id, seq_len) tuples.
        compute_estimator (Callable[[int, Optional[int]], float]): workload estimator function.
        max_seq_len_per_rank (int): max tokens per rank for packing.
        min_cp_size (int, optional): minimum CP size for dynamic CP.
    """
    seqlens = [seq_len for _, seq_len in sample_seqlens]
    unique_cp_sizes = {dcp_gpus_needed(seq_len, max_seq_len_per_rank, min_cp_size) for seq_len in seqlens}
    target_bucket_count = len(unique_cp_sizes)

    per_sample_work = []
    for _, seq_len in sample_seqlens:
        cp_size = dcp_gpus_needed(seq_len, max_seq_len_per_rank, min_cp_size)
        per_sample_work.append(compute_estimator(seq_len, cp_size))

    total_work = sum(per_sample_work)
    target_work_per_bucket = total_work / target_bucket_count

    buckets, current_bucket, current_bucket_work = [], [], 0.0
    remaining_target_buckets = target_bucket_count

    for sample_idx, (sample_id, seq_len) in enumerate(sample_seqlens):
        sample_work = compute_estimator(seq_len)
        projected_bucket_work = current_bucket_work + sample_work
        need_reserve_buckets_for_remaining_samples = len(sample_seqlens) - sample_idx <= (
            remaining_target_buckets - len(buckets)
        )
        exceeds_bucket_work_target = projected_bucket_work > (
            target_work_per_bucket * bucket_overfill_factor
        )
        should_close_current_bucket = bool(current_bucket) and (
            exceeds_bucket_work_target or need_reserve_buckets_for_remaining_samples
        )

        if should_close_current_bucket:
            buckets.append(deque(current_bucket))
            current_bucket, current_bucket_work = [], 0.0
            remaining_target_buckets -= 1

        current_bucket.append((sample_id, seq_len))
        current_bucket_work += sample_work

    if current_bucket:
        buckets.append(deque(current_bucket))
    return buckets

Some conventions for better readability:

1. Use meaningful name with clarity.
2. Simply complex condition
3. Prevent hardcoded magic number like 1.1 here.
4. Use consistent naming. In this PR, workload should be preferred than work/compute.

It seems this function doesn't guard the number of buckets equal to the number of unique CP sizes?
If so, the comments should be updated.

[CRITICAL Correctness] fill_empty inner while loop condition order causes Python negative-index wrap-around.

When i = 0, the loop evaluates sample_id_group[0] (true), then decrements i to -1. On the next iteration, Python evaluates sample_id_group[-1] (wrapping to the last element) before checking i >= 0. This inflates cp_size, causing align_sample_id_groups to split micro-batches incorrectly under VPP.

Suggestion:

# Before:
while sid0 in sample_id_group[i] and i >= 0:
# After (short-circuit prevents negative index access):
while i >= 0 and sid0 in sample_id_group[i]:

[CRITICAL Correctness] Assertion logic is inverted — allows overwriting non-empty GPU slots.

not all(work for ...) passes when at least one slot is empty, but the intent is "all slots in this range must be empty". With the current logic, if only 1 of needed_count slots is empty and the others contain real data, the assertion passes and the expansion overwrites existing assignments silently.

Suggestion:

# Before:
assert not all(
    work for work in micro_batches[empty_gpu : empty_gpu + needed_count]
), "Empty GPUs were detected but not enough to expand."
# After:
assert all(
    not work for work in micro_batches[empty_gpu : empty_gpu + needed_count]
), "Empty GPUs were detected but not enough contiguous empty slots to expand."

[IMPORTANT Correctness] Assert message is silently truncated — the second string line is a standalone expression, not part of the message.

Without enclosing parentheses, the newline after "... cannot redistribute, " terminates the assert statement. "try to increase 'max-seqlen-per-dp-cp-rank'." on the next line becomes a no-op expression statement.

Suggestion:

assert existing_group_sizes, (
    "There should be at least one group existing, cannot redistribute, "
    "try to increase 'max-seqlen-per-dp-cp-rank'."
)

[IMPORTANT Readability] next_hdp_group is 270 lines with 3 nested closures (trim_overload, fill_empty_gpus, inner fill_empty), mixing greedy scheduling, balance checking, empty-GPU filling, and overload trimming in a single scope.

The closures capture mutable outer state via nonlocal, making it hard to reason about which variables are modified where. fill_empty_gpus alone is ~70 lines of array-shifting logic that would be much easier to test and review as a standalone module-level function.

Suggestion: Extract fill_empty_gpus (and trim_overload when uncommented) into top-level functions with explicit input/output signatures. next_hdp_group should only contain the greedy main loop.

[SUGGESTION Naming] Core scheduling state variables are overly abbreviated for a 270-line function in distributed scheduling code.

In a function this long with multiple nested scopes, short names like gid, g_members, g_size, nc, next_pw force readers to keep a mental lookup table. Suggested renames:

[SUGGESTION Readability] Prefer isinstance(batch['local_cp_size'], int) over type(...) == int.

isinstance is the Pythonic idiom and correctly handles int subclasses (e.g. np.int64). Same applies to the type(batch['max_seqlen']) == int check on line 595.

[SUGGESTION Simplification] strategy parameter is always "dp" — the "pp" branch and pp_cursor variable are dead code.

No caller passes strategy="pp", so the round-robin scan_order path is unreachable. The pp_cursor maintenance (lines 675-676, 732) and the strategy parameter add complexity without providing value.

If "pp" is planned future work, add a # TODO with a tracking issue. Otherwise, remove the strategy parameter and the pp-related code to simplify the already complex scheduling logic.