Memory Profiling

[Kirankumarbs]

How do you debug OOM on long-running streaming jobs?

Hi All,

One of my Structured Streaming jobs has started dying with OutOfMemoryError. Runs fine for about approx 40 hours, then falls over, restarts, and does the same thing again. I had introduced few changes between stable and failed runs but changes are not obvious(Of course, i started slicing changes and testing them in isolation).

I've poked at it for a couple of days now. Spark UI first, obviously. Checked the Storage tab in case something cached was growing, looked at heap usage in the compute metrics to see whether it's a slow climb or a sudden spike (it's a climb), and read the executor logs around the crash. GC was getting progressively worse before it went.

My main suspect right now is leaking at Native Memory not Heap or Stateful Store etc

Where I'm stuck is knowing when to stop looking and just give it more memory. I could bump memoryOverhead tomorrow and the job would probably survive another week!!! How do you tell the difference? Is there a signal you look for that says `this is genuinely a leak` vs `this job was always borderline`?

Also curious whether anyone bothers with heap dumps on Databricks. Every time I've considered it the extraction and analysis looked like more effort than it was worth, but maybe I'm wrong about that.

And for the streaming people: what actually bites you most often? State growth, skew, someone leaving a collect() in a foreachBatch? I'd like to know what to check first next time instead of starting from scratch.

Thanks,
Kiran Kumar

[suranga]

Hi Kiran,

A heap that climbs steadily for 40 hours while GC gets worse looks like something holding on to memory, not a job that was always borderline. More memory would mostly move the crash from 40 hours to 80.

How I tell the two apart:

- Watch the floor, not the peaks. Look at heap used right after a full GC. If that floor keeps rising at a steady input rate, something is retaining memory. If it stays flat and the job only dies on spikes (a big batch, a backlog after a restart, a skewed key), it's borderline, and more memory or smaller batches is a fair fix.

- Heap or native. In the compute metrics, switch from All nodes to the driver and to single workers, and compare JVM heap usage with Container memory usage. If the container climbs while the heap stays flat, it's native memory (RocksDB, Python workers, Arrow). A java.lang.OutOfMemoryError means heap; a process killed by the OS (often exit code 137) points to native. GC getting worse points at the heap.

What I check first in streaming jobs:

1. State that never expires. In query.lastProgress, look at stateOperators: numRowsTotal and memoryUsedBytes. If numRowsTotal only goes up, an aggregation, dropDuplicates or stream-stream join has no watermark, or one much longer than it needs.

2. Where the state lives. Below DBR 17.3 LTS the default state store keeps state on the executor heap, which fits a slow heap climb. RocksDB moves it off-heap; on those versions also cap it with spark.sql.streaming.stateStore.rocksdb.boundedMemoryUsage and maxMemoryUsageMB. From 17.3 LTS, RocksDB is the default and its memory is capped automatically.

3. foreachBatch code. collect() or toPandas(), DataFrames cached and never unpersisted, or a Python list or dict that grows every batch. These grow the driver, not the executors.

Heap dumps are worth it once you know it's the heap, but start with something cheaper. On classic compute, take class histograms of the driver an hour apart from a %sh cell (jcmd <pid> GC.class_histogram) and compare the top classes. The one that keeps growing usually names the culprit. For a full dump, add -XX:+HeapDumpOnOutOfMemoryError -XX:HeapDumpPath=/local_disk0 to spark.driver.extraJavaOptions or spark.executor.extraJavaOptions, copy the .hprof to a volume, and open it in Eclipse MAT. It will be as big as the heap.

Docs:

https://docs.databricks.com/aws/en/structured-streaming/rocksdb-state-store

https://docs.databricks.com/aws/en/compute/cluster-metrics

// Suranga

[Kirankumarbs]

Thanks for reply and let me share my investigation so far!

> Look at heap used right after a full GC. If that floor keeps rising at a steady input rate, something is retaining memory.

Heap seems not raising. Live heap after GC is around 2-3.5GB(we are using 16GB machine D4s_v3) and heap is set by Databricks 7.67 -xmx and Init. For example the heap visualization.

Before GC:

 After GC:

 -----

> If the container climbs while the heap stays flat, it's native memory (RocksDB, Python workers, Arrow).

that's what error is and it is unclear from logs and also already checked that container memory is set by Databricks for this cluster is around 10.83 GB and we can't say container memory is raising but that's what all process RSS is used for! for example

In fact, it seems flat and we are using Single node(without workers)

-----

State Store also looks fine and from logs it shows live foot print is around 50MB and pinned memory by RocksDB is around 600MB and not growing! From Spark UI we can confirm this as well

  -----

for each batch collect() and persist and upersist looks good and does seem neither leak nor bloated complexity!

----

In fact Claude Code also suggested similar and tried but long long session with Claude code also not useful or not Productive. Hence, i am looking for first hand experience investigating memory issues beyond Heap!

As of now i am enabling below configs for the job to capture more information over process level metrics and OS level metrics:

              "spark.executor.processTreeMetrics.enabled": "true"
              "spark.executor.metrics.pollingInterval": "30s"
              "spark.eventLog.logStageExecutorMetrics": "true"

 I will draw few more visualization and see if that of any help!

[NityaOne]

Seems there are not many practical docs and videos around Databricks Jobs memory profiling!

[suranga]

Hi Kiran,

That's a really thorough round of checks. A flat heap floor, flat RocksDB and a small state store rule out the usual suspects, which narrows things down a lot.

The one thing I'd grab next is the exact OutOfMemoryError line from the driver log, because the text after the colon tells you which pool ran out:

• Java heap space or GC overhead limit exceeded: still the heap. With a flat floor that points to a sudden spike rather than a leak, like a big batch after a backlog or a collect/toPandas in foreachBatch.
• Direct buffer memory: off-heap NIO buffers, often from network clients or Arrow.
• unable to create native thread: a thread leak. Every thread costs native stack memory, which fits "heap flat, process slowly dies".
• Metaspace: classes being loaded over and over.

The metrics you've just switched on should confirm which one. With processTreeMetrics enabled, the executor metrics split RSS into ProcessTreeJVMRSSMemory, ProcessTreePythonRSSMemory and ProcessTreeOtherRSSMemory, and DirectPoolMemory and JVMOffHeapMemory are reported as well. With logStageExecutorMetrics on, the per-stage peaks land in the event log. On a single node the driver is also the executor, so it's all one JVM, which makes the numbers easier to read.

A cheap way to rule a thread leak in or out is to log the thread count every batch:

def log_threads(batch_df, batch_id):
    tm = batch_df.sparkSession._jvm.java.lang.management.ManagementFactory.getThreadMXBean()
    print(batch_id, tm.getThreadCount(), tm.getPeakThreadCount())
    # ... your existing batch logic

If that count keeps climbing over the 40 hours, look for anything created per batch and never closed: HTTP or database clients, thread pools, connection pools. _jvm access works on a dedicated (single user) cluster, not on standard access mode.

I'd be curious what the OOM line says once you have it.

Cheers,
Suranga