Skip to content

feat: add optional container memory guard based on cgroup usage - #5993

Closed
andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:feat/cgroup-memory-guard
Closed

andygrove wants to merge 1 commit into
apache:mainfrom
andygrove:feat/cgroup-memory-guard

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

Relates to #4576. Draft, opened for discussion rather than merge, and deliberately overlapping with @comphead's #5666 so the two approaches can be compared.

Rationale for this change

When an executor exceeds its container memory limit the kernel kills the whole JVM, taking every task on it, its cached blocks and its shuffle files. A single failed task is retried by Spark and costs almost nothing by comparison. This adds an optional guard that samples the container's real memory usage and fails the current task when usage approaches the limit.

The design differs from the two previous attempts in one respect, and it is the point of the PR: it reads the kernel's number rather than Comet's own accounting.

The closing review of #4582 set this precondition:

Whether to enforce on that number at all is a decision I'd rather make after seeing how well it tracks RSS on real workloads, rather than before.

I measured that and posted the result on #4576. On TPC-H SF100 with the default allocator, across ~4500 samples pairing native_allocated with kernel RSS at the same trace anchor:

  • delta correlation 0.42 to 0.44
  • the balance never leads RSS: correlation within 0.01 of zero at 1, 2, 5 and 10 sample horizons
  • the gap between them wanders by 2 to 3 GB, while the balance itself only spans 0 to 2.1 GB

A threshold on the allocator balance cannot mean anything about the quantity the kernel kills on, because the noise in the offset exceeds the whole signal. That is the same conclusion the #4582 review reached analytically ("the tracked balance is layout bytes, not RSS ... that gap is unbounded and always in the dangerous direction"), now with numbers.

Reading the cgroup avoids the problem entirely: it is the number the OOM killer compares against the limit, and it already includes the JVM heap, Comet's native allocations, JVM-side Arrow, Spark's own off-heap and mapped files. In my measurements the JVM heap was the dominant term, so any signal that excludes it is measuring the wrong thing.

What changes are included in this PR?

  • native/core/src/execution/memory_guard.rs (new). Discovers the container's memory usage file and limit: cgroup v2 at the mount root (the Kubernetes case, where the pod's cgroup is namespaced), then cgroup v2 resolved from /proc/self/cgroup for non-containerised hosts, then cgroup v1 for older Kubernetes and YARN. check() samples usage, throttled to at most once per 100 ms, and reports a trip at a configurable fraction of the limit.
  • jni_api.rs calls it at the two existing execution checkpoints: the async batch_receiver path and the ScanExec busy-poll path's 100-poll checkpoint. A trip becomes DataFusionError::ResourcesExhausted, which reaches the JVM as CometNativeException, which CometExecIterator logs with the task id and rethrows, so Spark fails and retries one task.
  • spark.comet.exec.memoryGuard.enabled (default false) and spark.comet.exec.memoryGuard.threshold (default 0.9). Names match feat: experiment with RealUsagePool #5666 so the two are comparable.
  • Tuning guide section.

No allocator wrapper, no panic_any, no stamped-thread set. That sidesteps the four defects the #4582 review found: the two enforcement layers not actually layering, the stamped set covering tokio's blocking pool so panics escaped as JoinError, the LOCAL_DRIFT leak on thread exit, and the layout-bytes-versus-RSS gap.

Limitations, stated up front

  • Checks happen between batches. An operator that grows sharply within a single batch can still cross the limit before the guard observes it. This narrows the window, it does not close it. Closing it entirely means acting inside the allocator, which feat: prototype allocator-level OOM protection (OomGuard breaker + cooperative real-usage accounting) #4582 tried and which brings reentrancy and unwinding problems of its own.
  • It fails the task rather than spilling first. Spilling releases reserved bytes, and in measurements on feat: observe real native memory usage in the off-heap memory pools #5983 the overshoot lived in allocations that never reserved, so spill-then-retry ran into the same wall and only delayed the failure.
  • No per-task attribution. The cgroup counter is process-wide, so the task that gets failed is the one that happened to reach a checkpoint, not necessarily the one responsible.
  • Not validated in a container. See testing below. This is the main reason it is a draft.

How are these changes tested?

Unit tests for the parsing, which is the part that can be tested off Linux: the cgroup v2 path from /proc/self/cgroup including the 0::/ container form, a v1-only host producing no v2 path, max treated as no limit, usage parsing, and out-of-range thresholds disabling the guard.

What is not tested, and I would not merge it without this: the guard has never run in a container. discover() is gated on a runtime cfg!(target_os = "linux") check, so the /sys/fs/cgroup reads typecheck on macOS but never execute there, and I have no Kubernetes environment to hand. What it needs is a pod with a memory limit, a query that exceeds it, and confirmation that the task fails and is retried while the executor survives. If anyone can run that, I would value it more than any amount of further code review.

When an executor exceeds its container memory limit the kernel kills the
whole JVM, taking every task on it along with its cached blocks and
shuffle files. Failing a single task is retried by Spark and costs far
less.

Sample the container's memory usage from the cgroup at existing
execution checkpoints and fail the current task once usage reaches a
configurable fraction of the limit. Read the kernel's number rather than
Comet's own allocator accounting: measurement on TPC-H SF100 showed the
allocator balance does not lead RSS at any horizon and that the offset
between them wanders further than the balance's entire range, so a
threshold on it cannot say anything about what the OOM killer compares
against.

Discovery covers cgroup v2 at the mount root, which is the Kubernetes
case, then cgroup v2 resolved from /proc/self/cgroup, then cgroup v1 for
older Kubernetes and YARN. With no container limit the guard disables
itself, so this is a no-op outside containers and off Linux.

Add spark.comet.exec.memoryGuard.enabled, defaulting to false, and
spark.comet.exec.memoryGuard.threshold, defaulting to 0.9.
@github-actions github-actions Bot added the enhancement New feature or request label Sep 17, 2026
@andygrove

Copy link
Copy Markdown
Member Author

Closing this.

The approach was validated end to end on k8s and the plumbing works: cgroup v2 discovery inside a pod, correct limit, trip propagating as CometNativeException. But memory.current turns out to be the wrong signal. It counts reclaimable page cache, so it saturates at the container limit on any workload that reads files.

In an A/B on TPC-H SF100 Q9 at a 9Gi pod limit, the guard failed the job (24 trips) while the identical run with the guard off completed in 58.62s, and nothing was OOMKilled in either. Measured anon peaked at 55% of the limit while memory.current sat at 99.99%.

Details and the cgroup trace are in #4576. A future attempt should threshold on anon/PSI rather than iterate on this branch.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant