[BUG]: Cache workers with no assigned data files never become ready, so the TrainJob hangs in its init container
- Dominant language
- Go
- Stars
- 2.2k
- Forks
- 1.1k
- Avg merge
- 3d 22h
- Merged PRs (30d)
- 39
Description
### What happened?
`partition_tasks` (`pkg/data_cache/src/head/provider.rs:377`) spreads an Iceberg table's data files across `num_workers` groups. When the table has fewer data files than the cache has workers, the surplus groups stay empty — and those workers are never told anything. They never become ready, so the TrainJob hangs in its init container with no error.
`cluster_size` defaults to `"3"` (`pkg/initializers/types/types.py:72`) and the head derives its worker list from `for i in 1..lws_size` (`pkg/data_cache/cmd/head/main.rs:47`), so a default install has two workers. An Iceberg table with a single data file is enough to hit this.
The chain:
1. `partition_tasks` leaves the surplus group with no tasks.
2. `RecordBatchBuilder::add_task` returns early for a group with no tasks (`provider.rs:302-305`), so that worker gets no row in the head's assignment table and `DataFileTableExec::execute` emits a zero-row batch for its partition.
3. `send_record_batch` routes each assignment by the `worker_ids` column of the batch carrying it (`pkg/data_cache/src/head/writer.rs:225-252`). A worker with no row is never a destination, so it never receives a `do_put`.
4. A worker registers its `memtable` only while handling a `do_put` (`pkg/data_cache/src/worker/worker_service.rs:329-331`), so the table never exists.
5. `check_table_ready` fails on `table_exist` (`pkg/data_cache/src/health/mod.rs:80-88`) and `/ready` returns 503 for the lifetime of the pod.
6. The LeaderWorkerSet never reports `Available`, and the dataset initializer's poll loop (`pkg/initializers/dataset/cache.py:311-333`) waits on that condition forever.
Two smaller defects sit in the same path:
- **An empty group gets an inverted row range.** `end_index` is computed as `end - 1` where `end` equals `start` (`provider.rs:421-427`), leaving `start_index` one above `end_index`. Nothing consumes that range today, because the row never reaches the assignment table, but it is not a range a consumer could reason about.
- **The zero-row batch from step 2 still reaches `send_record_batch`**, where `values.first()` fails with `No worker ID found` (`writer.rs:239-241`). That error aborts the partition's stream, so any batch queued behind it in the same partition is never delivered.
### Reproduction
Against `master` @ `8ca43ec`, two data files over three workers:
```
worker 0: tasks=1 start_index=0 end_index=99
worker 1: tasks=1 start_index=100 end_index=149
worker 2: tasks=0 start_index=150 end_index=149 <- start_index > end_index
DataFileTableExec::execute
partition 0 emitted 1 row(s)
partition 1 emitted 1 row(s)
partition 2 emitted 0 row(s) <- worker 2 is never a do_put destination
```
The existing tests do not reach this case. `test_partition_tasks_validates_start_end_indices_empty` (`provider.rs:535`) uses an all-empty stream, where every group gets `(0, 0)` and no inversion can appear. `test_partition_tasks_validates_start_end_indices_varying_groups` (`provider.rs:572`) skips its `end_index` assertion whenever a group's record count is 0 (`provider.rs:634`). The mixed case — some groups populated, some empty — is untested.
### What did you expect to happen?
A cache cluster should serve a table that has fewer data files than it has workers. A worker that legitimately holds no rows should still report ready so the LeaderWorkerSet becomes available, and an empty group should carry a range a consumer can reason about rather than an inverted one.
### Environment
Reproduced from source at `master` @ `8ca43ec` (VERSION `v2.3.0`) by running `cargo test` against `pkg/data_cache`. The output above is unit-level; I have not run it against a live cluster.
Kubernetes version:
```bash
$ kubectl version
n/a - reproduced from source
```
Kubeflow Trainer version:
```bash
$ kubectl get pods -n kubeflow-system -l app.kubernetes.io/name=kubeflow-trainer -o jsonpath="{.items[*].spec.containers[*].image}"
n/a - reproduced from source at master @ 8ca43ec (VERSION v2.3.0)
```
Kubeflow Python SDK version:
```bash
$ pip show kubeflow
n/a
```
### Related
#3927 covers the initializer's unbounded wait for the LeaderWorkerSet. This is one reason that wait never terminates, so the two are complementary rather than duplicates.
### Impacted by this bug?
Give it a 👍 We prioritize the issues with most 👍
Contributor guide
Assessment
This issue has not been assessed yet.