fix: Fairly allocate partitions from scheduler assignment cycle only

[peasee]

📝 Summary

• Updates the scheduler to return Unavailable to executors requesting their initial partition allocations until the first allocation assignment interval passes. This ensures that executors have some time to connect - instead of the first executor racing to greedily allocate all of the partitions for a table to itself
• Updates initial allocation to use fair share allocation similar to the regular partition interval assignments, instead of greedily assigning every partition for a table to one executor
• Ensures that partitions are more evenly spread across executors on startup to prevent executors taking too many

👀 Notes for Reviewers

• Partition assignment probably needs more work. Returning Unavailable to the executors forces them to re-check their initial assignments in a retry loop, but if the partition assignment interval is too long that retry loop will fail before the scheduler allocates any initial partitions. I think partition assignment, including initial partitions, needs to be entirely push-based instead of executors pulling partitions on first startup - but that requires some additional work.

[github-actions]

✅ Pull with Spice Passed

Passing checks:

• ✅ Title meets minimum length requirement (10 characters)
• ✅ No banned labels detected
• ✅ Has a label from required category kind/
• ✅ Has a label from required category area/
• ✅ Has at least one assignee: peasee

[Copilot]

Pull request overview

This PR refactors cluster partition assignment so initial partition ownership is determined by the scheduler’s periodic fair-share assignment cycle (instead of executors greedily allocating on connect), and makes executors retry allocate_initial_partitions until the scheduler is ready to return a fair initial share.

Changes:

• Delay the first scheduler partition-assignment cycle by one interval and introduce a “first assignment complete” gate to control when executors can fetch initial assignments.
• Change allocate_initial_partitions to return only partitions already assigned to the requesting executor (no longer mutating assignments).
• Add an integration test assertion to validate fair-share distribution across multiple executors at startup.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
crates/runtime/tests/cluster/harness.rs	Adds helper to read per-executor partition ownership from scheduler metadata for assertions.
crates/runtime/tests/cluster/distributed_acceleration.rs	Adds a test assertion that initial partitions are fairly split across two executors.
crates/runtime/src/cluster/service.rs	Updates allocate_initial_partitions to gate on first assignment and stop performing allocation.
crates/runtime/src/cluster/partition/scheduler_task.rs	Delays first assignment tick; marks the “first assignment complete” gate after a successful cycle.
crates/runtime-cluster/src/store.rs	Removes executor-driven allocate_partitions API and the associated AllocationResult type.
crates/runtime-cluster/src/service.rs	Adds first_assignment_complete gate state to the partition service.
crates/runtime-cluster/src/lib.rs	Removes re-export of AllocationResult after its deletion.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

crates/runtime/src/cluster/partition/scheduler_task.rs:124

• The first assignment cycle is deferred by self.interval, but executors only retry allocate_initial_partitions for a bounded time on Unavailable (see ALLOCATE_INITIAL_PARTITIONS_MAX_RETRIES + SCHEDULER_BACKOFF_MAX in crates/runtime/src/cluster/mod.rs:86-92, ~5 minutes worst-case). If partition_assignment_interval is configured larger than that, executors can hard-fail startup before the scheduler ever runs its first cycle and opens the gate.

            delay_ms = self.interval.as_millis(),
            "Deferring first partition assignment cycle to let executors connect"
        );
        let start = tokio::time::Instant::now() + self.interval;
        let mut interval = tokio::time::interval_at(start, self.interval);

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

crates/runtime/src/cluster/partition/scheduler_task.rs:123

• The first partition-assignment cycle is deferred by a full partition_assignment_interval. Since executors retry allocate_initial_partitions for a bounded time (~5 minutes total), configuring a long interval (e.g. >5m) will reliably prevent executors from starting because the scheduler won't assign anything before the retry budget is exhausted. Consider using a small bounded startup window (or a dedicated config field) for the initial delay instead of coupling it to the steady-state interval.

        tracing::info!(
            delay_ms = self.interval.as_millis(),
            "Deferring first partition assignment cycle to let executors connect"
        );
        let start = tokio::time::Instant::now() + self.interval;

crates/runtime/src/cluster/service.rs:699

• The first-assignment gate blocks all allocate_initial_partitions calls until the scheduler runs an assignment pass. On scheduler restart, the partition store may already contain assignments for an executor, but this gate will still force the executor to retry (and potentially time out) even though the scheduler could safely return the already-assigned partitions immediately. Consider bypassing the gate when the cached partition metadata shows this executor already owns at least one partition.

        if let Some(partition_service) = self.datafusion.partition_service.as_ref()
            && !partition_service.is_first_assignment_complete()
            && self.has_partitioned_accelerated_tables().await
        {
            tracing::debug!(

[phillipleblanc]

I think the correct solution would be to have a proper rebalancing control loop that doesn't depend on executors coming in right after startup, but this at least is an incrementally better solution than what we already had.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

crates/runtime/src/cluster/partition/scheduler_task.rs:116

• Deferring the first partition-assignment cycle by a full self.interval means executors will receive Unavailable from allocate_initial_partitions for at least one whole interval (and potentially nearly two intervals if they connect just after the first tick and ReconcileOutcome::NoAssignment keeps the gate closed).

With long partition_assignment_interval values this can cause executor startup to fail after exhausting its Unavailable retry budget (see crates/runtime/src/cluster/mod.rs: SCHEDULER_BACKOFF_MAX + ALLOCATE_INITIAL_PARTITIONS_MAX_RETRIES). Consider decoupling the startup grace period from the steady-state assignment interval (e.g., separate config, or a bounded grace window) so “fair startup” doesn’t become “minutes-long startup stall”.

        // Defer the *first* assignment cycle by one full interval so executors
        // have a window to connect before the scheduler distributes partitions.
        // Assignment is scheduler-controlled (executors no longer allocate on
        // connect via `allocate_initial_partitions`), so an executor that joins
        // late — but within this window — is included in the initial fair