apache / apache/datafusion-ballista

Cancelling a running job leaks its executor vcores and wedges the scheduler under PushStaged

Open
#2,418 0 comments 1 reaction 0 assignees View on GitHub
bug
Dominant language
Rust
Stars
2.1k
Forks
320
Avg merge
1d 22h
Merged PRs (30d)
66

Description

## Describe the bug

Cancelling a running job never returns the vcores its tasks had reserved. Under the default `PushStaged` scheduling policy the scheduler's view of the executor's free vcores is permanently reduced by the cancelled tasks' consumption, so after enough cancellations the scheduler has nothing left to bind and **every subsequent query hangs forever**. Registering a fresh executor does not help.

With a 16-vcore executor and a single collapse task that monopolizes it, one cancellation is enough to wedge the cluster.

## To Reproduce

Against `main` (verified at d9fadd8c9), with a Parquet TPC-H dataset:

```sh
./target/release/ballista-scheduler & # default PushStaged
./target/release/ballista-executor -c 16 -p 50051 &
```

`short.sql`:
```sql
CREATE EXTERNAL TABLE lineitem STORED AS PARQUET LOCATION '/path/to/tpch/lineitem';
SELECT count(*) FROM lineitem;
```

`long.sql`: same DDL, plus a query that runs for a while:
```sql
SELECT count(*) FROM lineitem l1 JOIN lineitem l2
ON l1.l_orderkey = l2.l_orderkey AND l1.l_linenumber = l2.l_linenumber;
```

1. `ballista-cli --host localhost --port 50050 -f short.sql` — returns in ~0.07s.
2. Start `ballista-cli --host localhost --port 50050 -f long.sql` in the background.
3. Once it is running, cancel it through the scheduler's own API: `SchedulerGrpc/CancelJob { job_id }` (the job id is in the scheduler's `Job submitted: [...]` log line). The client correctly reports `Job failed: Cancelled`.
4. Run `short.sql` again. It never returns.

Scheduler log with `RUST_LOG=ballista_scheduler=debug`:

```
INFO ballista_scheduler::state::task_manager: Cancelling 1 running tasks for job 038392DNSJ000
ERROR ballista_scheduler::state::task_manager: Fail to find job 038392DNSJ000 in the active cache and it may not be curated by this scheduler
DEBUG ballista_scheduler::cluster: No executor vcores available for task binding
DEBUG ballista_scheduler::state: No schedulable tasks found to be launched
```

The third line is the symptom: the executor is alive and idle, but the scheduler believes it has no capacity.

## Expected behavior

Cancelling a job returns the vcores its running tasks held, and the cluster keeps accepting work.

## Additional context

Under `PushStaged`, vcores are refunded in exactly one place — the `TaskUpdating` handler in `scheduler_server/query_stage_scheduler.rs`, which calls `TaskManager::sum_vcores_for_statuses` (`state/task_manager.rs:515`). That function resolves each status through `get_active_execution_graph`, i.e. the **active job cache**.

On cancellation the order is:

1. `JobCancel` -> `abort_job` marks the stages failed, asks the executor to cancel its tasks, and then `persist_terminal_and_evict` removes the job from the active cache.
2. The executor reports the cancelled task's status afterwards.
3. `sum_vcores_for_statuses` cannot find the job (hence the `Fail to find job ... in the active cache` error above), returns 0, and nothing is refunded.

The `PullStaged` path does not have this problem: `poll_work` in `scheduler_server/grpc.rs` fails the job on a preparation error and does not use this budget accounting at all.

Two things make a fix worth some care:

- `sum_vcores_for_statuses` takes the graph's **read** lock and `update_task_statuses` takes the **write** lock separately, so a cancellation interleaving between them could double-count a task if the abort path also tallies it.
- `InMemoryClusterState::unbind_tasks` (`cluster/memory.rs:113`) does `data.vcores += num_vcores` with no clamp against the executor's `total_vcores`, so an over-refund would silently oversubscribe an executor rather than fail loudly. That is arguably worth fixing on its own regardless.

Found while testing #2416; it is unrelated to that PR and reproduces on `main` with no Flight SQL involved. It does mean a Flight SQL frontend that cancels abandoned queries cannot land until this is fixed, since it would turn a wasted query into a wedged scheduler.

Contributor guide

Open the contributing guide

Research direction

Reproduce the cancellation sequence with PushStaged, then read scheduler_server/query_stage_scheduler.rs, state/task_manager.rs, and cluster/memory.rs, focusing on TaskUpdating, sum_vcores_for_statuses, and unbind_tasks. Trace JobCancel through abort_job and persist_terminal_and_evict, and inspect the PullStaged path for comparison. Done means cancelling a running task refunds its vcores exactly once and a subsequent query completes without oversubscribing the executor.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
65/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.