[andygrove]

Which issue does this PR close?

No single issue. Split out of #2792. Related: #2520 (performance review tracking).

Rationale for this change

Under pull scheduling the scheduler binds tasks, marking them running, before tonic encodes the PollWorkResult. When the response exceeds the gRPC encoding limit (many tasks with large plans, such as broadcast readers listing every upstream location at 128 partitions), the executor gets an error, the tasks never run, and the job hangs.

What changes are included in this PR?

The scheduler's poll handler keeps the tasks that fit under the gRPC encoding limit (--grpc-server-max-encoding-message-size), counting the cleanup list in the same response, and returns the rest to their stages. A task that does not fit does not stop smaller ones behind it from being kept.

• Returned tasks are reported to the stage as FailedTask::from(BallistaError::Cancelled), the retryable TaskKilled status executors send for a cancelled task. The stage re-queues their partitions, and the tasks are bound again on a later poll.
• A task larger than the limit by itself can never be sent, so its job fails with an error naming the setting to raise (and the executor's decoding limit).

Unit tests check the size estimate against prost's real encoding (a limit of exactly three tasks plus a cleanup entry keeps three, one byte less keeps two) and that small tasks behind an oversized one are kept.

Known limits:

• The tasks are prepared (plan restricted and encoded) before their size is known, so a returned task pays that cost again when it is rebound. Bounding the number of tasks before preparing them would need the size in the bind functions, which hold the graph lock and have no codec.
• Each return bumps the partition's attempt counter, as every retryable failure does, including the other TaskKilled paths. The counter is also the attempt number used in shuffle file names, so it must keep increasing. It is compared with max_task_failures only when a task fails with count_to_failures.
• fit_tasks_in_response itself has no unit test, since it needs a running job; the splitting is tested.
• The limit checked is the scheduler's encoding limit. An executor with a smaller decoding limit is not detected.

Are there any user-facing changes?

No config or API changes. A job that previously hung now runs, or fails with an error that names the gRPC message-size setting.

[milenkovicm]

Thanks @andygrove, I have observed this issue few times, good fix. Making jobs fail instead of hanging is a good fix

Ballista plan is encoding a bit too much information which #1851 outlines, we should try to find a way to reduce it