khanlab / khanlab/SPIMquant

vessel_graph_to_nodes_edges: dask P2P shuffle OOMs on large vessel graphs

Open
#193 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
8
Forks
3
PR merge metrics
No merged PRs in 30d

Description

The dask-dataframe merge in `convert_vessel_graph_to_nodes_edges.py` fails on large vessel graphs. Workers are killed on memory pressure, which drops participants from an in-flight P2P shuffle and takes the whole merge down.

## Observed

3-subject run, `--executor slurm --default-resources -j 6`. One of three subjects failed:

```
Error in rule vessel_graph_to_nodes_edges:
jobid: 70
input: sub-AS11F3/..._desc-vesselfm+skeleton_graph.parquet
output: sub-AS11F3/..._desc-vesselfm+skeleton_nodes.parquet,
sub-AS11F3/..._desc-vesselfm+skeleton_edges.parquet
```

From the job log:

```
Worker memory limit: 3.91 GiB # workers climbing to 2.0-2.4 GiB, then killed
P2PConsistencyError("Worker tcp://127.0.0.1:34139 left during active
SchedulerShuffleState<...[172]>")
```

Counts across that one log: 378 `P2PConsistencyError`, ~90 nanny restarts.

## Why it is marginal, not a hard limit

Input `graph.parquet` sizes for the three subjects:

| subject | graph.parquet | result |
|---|---|---|
| AS8M4 | 3528 MB | ok |
| AS17F6 | 4180 MB | ok |
| AS11F3 | 4198 MB | **failed** |

The failing case is only ~0.4% larger than one that passed, so this is not a clean threshold. It will surface intermittently on any cohort with graphs in this size range rather than failing predictably.

## Likely cause

`vessels.smk` requests `threads: 128` with `mem_mb=256000`. Dask splits that into roughly 64 workers at ~3.9 GiB each. The double merge at `convert_vessel_graph_to_nodes_edges.py:355-359`

```python
edges_with_nodes = edges_ddf.merge(...).merge(...)
```

triggers a P2P shuffle whose per-worker partitions do not fit in ~4 GiB. Total memory is not the constraint; per-worker memory is. Raising `mem_mb` alone would not help while the worker count scales with it.

## Suggested direction

Fewer, fatter workers for this rule: reduce `threads`, or configure the cluster with an explicit worker count so each gets substantially more than 4 GiB. Repartitioning the edge table before the merge, or setting an explicit index to avoid the shuffle, would be a more durable fix.

## Notes

- Reproduced on `plaques` at 4a28212, but this rule and script are byte-identical to `main` there. Unrelated to the LANTERN work on that branch.
- Environment: pixi `gpu` env, dask/distributed as pinned in `pixi.lock`.
- Everything else in the run completed: 233 of 235 steps.

Filed with assistance from Claude Code.

Contributor guide

Open the contributing guide

Research direction

Start with vessels.smk and convert_vessel_graph_to_nodes_edges.py, especially the double merge at lines 355-359, and reproduce the vessel_graph_to_nodes_edges rule on a large graph. Compare worker sizing and shuffle behavior against the reported 4 GiB-per-worker failure. Done means large vessel graphs complete without worker deaths or P2PConsistencyError failures and produce both nodes and edges outputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering, distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.