Ray OutOfMemoryError: Task was killed due to the node running low on memory

We use Ray on Kubernetes (KubeRay) to preprocess images before computing embeddings: decode, resize, normalise, and write to Parquet. The job ran for about an hour and then failed, with workers dying repeatedly:

ray.exceptions.OutOfMemoryError: Task was killed due to the node running low on memory.
Memory on the node (IP: 10.42.3.17, ID: 8c2f...) where the task (task ID: a41b..., name=preprocess_batch,
pid=3812, memory used=3.12GB) was running was 29.61GB / 30.89GB (0.958591), which exceeds the memory
usage threshold of 0.95. Ray killed this worker (ID: 5d0e...) because it was the most recently scheduled
task; to see more information about memory usage on this node, use `ray logs raylet.out -ip 10.42.3.17`.

Retries hit the same limit, and after the retry budget was exhausted the job failed.

Reproducing it

The Ray dashboard showed each worker node running 16 preprocess_batch tasks at once, one per CPU, with each task's memory growing to about 3 GB at peak. 16 times 3 GB is 48 GB on a node with 30 GB. The first part of the dataset contained small thumbnails; after an hour, the job reached a folder of 24-megapixel product photos.

What we ruled out

• A memory leak. Per-task memory returned to baseline after each batch.
• Object store spilling as the main cause. The object store was busy but spilling to disk worked; the kills were for worker heap memory.

Root cause

Ray schedules tasks by declared resources. Our task declared only num_cpus=1 (the default), so Ray placed as many tasks on a node as it had CPUs, without considering memory. Each task loaded a batch of 64 full-resolution images into memory as decoded arrays, about 3 GB for large photos. Ray's memory monitor protects the node from the Linux OOM killer by killing workers when node memory use passes 95 percent, which it did.

What is happening under the hood

Ray has two kinds of memory on each node:

• Worker heap memory: what your task's Python process allocates. Not limited by Ray unless you declare a memory resource, which Ray uses for scheduling (it does not enforce a hard limit per process).
• Object store memory: shared memory for objects passed between tasks, 30 percent of node memory by default, with spilling to disk when full.

The memory monitor checks node memory periodically and, above the threshold (RAY_memory_usage_threshold, default 0.95), kills a worker, preferring retriable tasks and the most recently scheduled ones. Without memory declarations, the scheduler has no idea that tasks are memory-heavy.

Environment

Minimal reproduction

import numpy as np, ray

@ray.remote
def heavy():
    a = np.ones((3 * 1024**3) // 8)      # ~3 GB
    return a.sum()

ray.get([heavy.remote() for _ in range(64)])   # many run concurrently per node -> OOM kills

The fix

1. Declare memory per task. We measured peak memory per batch for the largest images and declared it, so Ray schedules at most as many tasks as fit:

@ray.remote(num_cpus=1, memory=3.5 * 1024**3)
def preprocess_batch(paths: list[str]) -> str:
    ...

On a 30 GB node, this allows 8 concurrent tasks instead of 16, leaving room for the raylet and object store.

2. Smaller, size-aware batches. Batches are now built by total pixel count rather than a fixed number of images, so a batch of large photos contains fewer images:

def batches_by_pixels(items, max_pixels=400_000_000):
    batch, total = [], 0
    for path, w, h in items:
        if batch and total + w * h > max_pixels:
            yield batch
            batch, total = [], 0
        batch.append(path)
        total += w * h
    if batch:
        yield batch

Images are also resized as soon as they are decoded, so the full-resolution array is released early.

3. Right-size the object store. The job passes small file paths between tasks, not image data, so we lowered object store memory for these workers to leave more for heap:

# KubeRay worker group rayStartParams
rayStartParams:
  object-store-memory: "4000000000"

4. Move to Ray Data for the pipeline, which streams blocks through stages with backpressure and lets us set memory per operator. That was the longer-term change; the first three fixes made the existing job reliable.

Alternatives and trade-offs

• Bigger nodes. More memory per CPU solves it with money; memory-optimised instances were not available in our quota at the time.
• Raise the memory threshold. Moves the kill closer to the kernel OOM killer, which kills less selectively and can take down the raylet. We kept the default.
• Fewer CPUs per task declaration (for example num_cpus=2) as an indirect memory limit. Works, but wastes CPU when tasks are not memory-heavy.

How to test the fix

We ran the job against a sample with the 5,000 largest images first. Peak node memory stayed at 78 percent, no tasks were killed, and throughput per node was only 12 percent lower than the original configuration on small images.

Confirming in production

• The full 3.2M-image job completed in 5 hours 40 minutes with zero memory-monitor kills.
• Ray dashboard shows 8 concurrent preprocessing tasks per node and stable memory.

Regression guard

• Every Ray task and actor in our shared library must declare memory; a decorator wrapper enforces it.
• An alert on Ray's task failure metrics with OutOfMemoryError as the reason.
• A pre-flight sample run on the largest inputs for new datasets.

Security implications

None directly. Large image datasets include user-uploaded content; the workers run with read access to the source bucket and write access only to the output prefix, so a crash or misconfiguration cannot overwrite source data.

Prevention checklist

• Declare memory for memory-heavy Ray tasks and actors.
• Size batches by data volume, not item count.
• Release large objects early; resize or downsample immediately after decoding.
• Size the object store for what you actually pass between tasks.
• Test on the largest inputs before the full run.

References

• Ray out-of-memory prevention — https://docs.ray.io/en/latest/ray-core/scheduling/ray-oom-prevention.html
• Ray resources and memory scheduling — https://docs.ray.io/en/latest/ray-core/scheduling/resources.html
• Ray memory management and object store — https://docs.ray.io/en/latest/ray-core/scheduling/memory-management.html
• Ray Data execution and memory configuration — https://docs.ray.io/en/latest/data/data-internals.html