[SPARK-50768][CORE] Introduce TaskContext.createResourceUninterruptibly to avoid stream leak by task interruption #49413

Ngone51 · 2025-01-08T08:59:04Z

What changes were proposed in this pull request?

This PR fixes the potential stream leak issue by introduing TaskContext.createResourceUninterruptibly.

When a task is using TaskContext.createResourceUninterruptibly to create the resource, the task would be marked as uninterruptible. Thus, any interruption request during the call to TaskContext.createResourceUninterruptibly would be delayed until the creation finishes.

This PR introduces an new lock contention between Task.kill and TaskContext.createResourceUninterruptibly. But I think it is acceptable given that both are not on the hot-path.

(I will submmit a followup to apply TaskContext.createResourceUninterruptibly in the codebase if this PR is approved by the community.)

Why are the changes needed?

We had #48483 tried to fix the potential stream leak issue by task interruption. It mitigates the issue by using

def tryInitializeResource[R <: Closeable, T](createResource: => R)(initialize: R => T): T = {
  val resource = createResource
  try {
    initialize(resource)
  } catch {
    case e: Throwable =>
      resource.close()
      throw e
  }
}

But this utility function has an issue that resource.close() would leak open resouces if initialize(resource) also created some resources internally, especially when initialize(resource) is interrupted (See the example of InterruptionSensitiveInputStream in the test).

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Added a unit test.

Was this patch authored or co-authored using generative AI tooling?

No.

core/src/test/scala/org/apache/spark/JobCancellationSuite.scala

core/src/main/scala/org/apache/spark/scheduler/Task.scala

core/src/main/scala/org/apache/spark/TaskContextImpl.scala

xuanyuanking · 2025-01-11T00:46:18Z

core/src/test/scala/org/apache/spark/JobCancellationSuite.scala

+
+            // Leave some time for the task to be interrupted during the
+            // creation of `InterruptionSensitiveInputStream`.
+            Thread.sleep(5000)


How important is this sleep within the task? Could it potentially make the test flaky?

It is necessary to ensure the task is interrupted during InterruptionSensitiveInputStream#initialize() so that we can test the leaked stream. Increase the sleep time should less likely be flaky.

core/src/test/scala/org/apache/spark/JobCancellationSuite.scala

core/src/main/scala/org/apache/spark/TaskContext.scala

xuanyuanking

LGTM

Ngone51 · 2025-01-15T03:18:11Z

Could someone help merge this PR if we're good to proceed? Or if I can also merge myself with enough approvals?

LuciferYang · 2025-01-15T03:33:10Z

@Ngone51 Does branch-3.5 also need this bug fix?

Ngone51 · 2025-01-15T03:36:34Z

@LuciferYang Thanks. Yes, I think so. Do I need send a separate PR?

LuciferYang · 2025-01-15T03:37:14Z

Merged into master for Spark 4.0. Thanks @Ngone51 @xuanyuanking and @HyukjinKwon .

@Ngone51 Due to code conflicts, it cannot be directly merged into branch-3.5. If needed, please submit a separate pr. Thanks ~

Ngone51 · 2025-01-15T03:38:34Z

Sure, I'd submit a separate PR for branch-3.5.

Ngone51 · 2025-01-15T03:54:05Z

Hi @LuciferYang I found we didn't backport #48483 to branch-3.5 for the same issue. Generally speaking, this is not a critial bug fix but only for the corner case. Not sure if we want to backport it or not.

@HyukjinKwon Any idea about the backport policy for this kind of fix? is it a requirement or optional?

HyukjinKwon · 2025-01-15T04:41:38Z

please go ahead. I think it's fien

dongjoon-hyun · 2025-01-15T04:46:17Z

Please hold one backporting for now, @Ngone51 .

It's because this could introduce a flakiness into master branch. branch-3.5 can wait because we don't have Apache Spark 3.5.5 release schedule yet.

In addition, if you can, please give us a confirmation that the failures of current master branch are irrelevant to this PR.

dongjoon-hyun · 2025-01-15T04:48:57Z

For the record, after this PR, surprisingly, 3 commits fail consecutively at the same PySpark pipeline.

dongjoon-hyun · 2025-01-15T05:01:32Z

core/src/main/scala/org/apache/spark/BarrierTaskContext.scala

+  override private[spark] def interruptible(): Boolean = taskContext.interruptible()
+
+  override private[spark] def pendingInterrupt(threadToInterrupt: Option[Thread], reason: String)
+  : Unit = {


nit. Indentation?

https://github.com/databricks/scala-style-guide?tab=readme-ov-file#spacing-and-indentation

I'll fix it with a followup. Thanks!

Fixed at #49508

dongjoon-hyun · 2025-01-15T05:01:45Z

core/src/main/scala/org/apache/spark/BarrierTaskContext.scala

+  }
+
+  override private[spark] def createResourceUninterruptibly[T <: Closeable](resourceBuilder: => T)
+  : T = {


ditto. Indentation.

https://github.com/databricks/scala-style-guide?tab=readme-ov-file#spacing-and-indentation

HyukjinKwon · 2025-01-15T05:14:16Z

im taking a look now for the failure

dongjoon-hyun · 2025-01-15T05:15:30Z

Thank you so much! I looked at, but didn't find any clear clue.

LuciferYang · 2025-01-15T05:47:06Z

Seems unrelated, I tested the result of reverting the current one, but pyspark.conf still failed:

HyukjinKwon · 2025-01-15T05:47:12Z

found the cuase. working on the fix

HyukjinKwon · 2025-01-15T05:47:22Z

no tihs pr is fine

LuciferYang · 2025-01-15T05:51:36Z

@HyukjinKwon Are you saying this pr caused the test to fail?

HyukjinKwon · 2025-01-15T05:53:53Z

nono this PR is fine. I will make a PR soon with a PR description with 20 mins

HyukjinKwon · 2025-01-15T06:10:36Z

here: #49500

cloud-fan · 2025-01-15T06:59:59Z

core/src/main/scala/org/apache/spark/TaskContextImpl.scala

+  // Should be protected by `TaskContext.synchronized`.
+  private var pendingInterruptRequest: Option[(Option[Thread], String)] = None
+
+  // Whether this task is able to be interrupted. Should be protected by `TaskContext.synchronized`.


or maybe mark as @transient?

Are you referring to @volatile?

Since we already have TaskContext.synchronized protected, @volatile could be redundant. And if we use @volatile only, it's not thread safety as we could have two threads modify _interruptible concurrently.

cloud-fan · 2025-01-15T07:07:26Z

core/src/main/scala/org/apache/spark/TaskContextImpl.scala

+
+  override def interruptible(): Boolean = TaskContext.synchronized(_interruptible)
+
+  override def pendingInterrupt(threadToInterrupt: Option[Thread], reason: String): Unit = {


This API looks weird if threadToInterrupt is None, as there is nothing to interrupt.

When threadToInterrupt=None, it would still "interrupt" the task by invoking TaskContext.markInterrupted(). It just won't invoke Thread.interrupt() on the task thread.

cloud-fan · 2025-01-15T07:13:13Z

core/src/main/scala/org/apache/spark/TaskContextImpl.scala

+    }
+
+    TaskContext.synchronized {
+      interruptIfRequired()


Is it a safeguard that the caller may mistakenly call pendingInterrupt even if the task is interruptable?

No. This is for the case where resource creation happens after the task has been marked as interrupted. In that case, pendingInterruptRequest is None and reasonIfKilled is not None, and killTaskIfInterrupted() would throw TaskKilledException() to stop the task thread.

Ngone51 · 2025-01-15T14:24:56Z

FYI I created a followup PR (#49508) to use TaskContext.createResourceUninterruptibly() where it applies.

…ninterruptibly() to risky resource creations ### What changes were proposed in this pull request? This is a follow-up PR for #49413. This PR intends to apply `TaskContext.createResourceUninterruptibly()` to the resource creation where it has the potential risk of resource leak in the case of task cancellation. ### Why are the changes needed? Avoid resource leak. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? n/a ### Was this patch authored or co-authored using generative AI tooling? No. Closes #49508 from Ngone51/SPARK-50768-followup. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]>

…ninterruptibly() to risky resource creations ### What changes were proposed in this pull request? This is a follow-up PR for #49413. This PR intends to apply `TaskContext.createResourceUninterruptibly()` to the resource creation where it has the potential risk of resource leak in the case of task cancellation. ### Why are the changes needed? Avoid resource leak. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? n/a ### Was this patch authored or co-authored using generative AI tooling? No. Closes #49508 from Ngone51/SPARK-50768-followup. Authored-by: Yi Wu <[email protected]> Signed-off-by: Dongjoon Hyun <[email protected]> (cherry picked from commit 9b32334) Signed-off-by: Dongjoon Hyun <[email protected]>

…atabricks] (#11972) This fixes #11971 . This unblocks the scala_2.13 nightly build jobs. MockTaskContext extends TaskContext in spark-rapids code. Recently additional definitions were added in Spark-4.0's TaskContext. In this PR, we added the missing methods to MockTaskContext.scala file. We are not overriding these methods so that it can compile on all Spark versions. In Spark-4.0, these were added in TaskContext.scala: PR : apache/spark#49413 ``` private[spark] def interruptible(): Boolean private[spark] def pendingInterrupt(threadToInterrupt: Option[Thread], reason: String): Unit private[spark] def createResourceUninterruptibly[T <: Closeable](resourceBuilder: => T): T ``` Signed-off-by: Niranjan Artal <[email protected]> Co-authored-by: Gera Shegalov <[email protected]>

fix

a962e83

github-actions bot added the CORE label Jan 8, 2025

Ngone51 requested review from mridulm, cloud-fan, jiangxb1987, HyukjinKwon and LuciferYang January 8, 2025 08:59

LuciferYang reviewed Jan 8, 2025

View reviewed changes

core/src/test/scala/org/apache/spark/JobCancellationSuite.scala Outdated Show resolved Hide resolved

core/src/main/scala/org/apache/spark/scheduler/Task.scala Outdated Show resolved Hide resolved

core/src/main/scala/org/apache/spark/TaskContextImpl.scala Outdated Show resolved Hide resolved

address Lucifer's comments

721e294

xuanyuanking reviewed Jan 11, 2025

View reviewed changes

Ngone51 added 3 commits January 13, 2025 23:04

fix mima

cd5b8db

fix fmt

eff0717

update comments

3eb2331

github-actions bot added the BUILD label Jan 13, 2025

increase sleep time

7f36c5b

xuanyuanking approved these changes Jan 13, 2025

View reviewed changes

LuciferYang approved these changes Jan 14, 2025

View reviewed changes

HyukjinKwon approved these changes Jan 14, 2025

View reviewed changes

LuciferYang closed this in 6f3b778 Jan 15, 2025

dongjoon-hyun reviewed Jan 15, 2025

View reviewed changes

cloud-fan reviewed Jan 15, 2025

View reviewed changes

Ngone51 mentioned this pull request Jan 15, 2025

[SPARK-50768][SQL][CORE][FOLLOW-UP] Apply TaskContext.createResourceUninterruptibly() to risky resource creations #49508

Closed

nartal1 mentioned this pull request Jan 16, 2025

Update MockTaskContext to support new functions added in Spark-4.0 [databricks] NVIDIA/spark-rapids#11972

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-50768][CORE] Introduce TaskContext.createResourceUninterruptibly to avoid stream leak by task interruption #49413

[SPARK-50768][CORE] Introduce TaskContext.createResourceUninterruptibly to avoid stream leak by task interruption #49413

Ngone51 commented Jan 8, 2025 •

edited

Loading

xuanyuanking Jan 11, 2025

LuciferYang Jan 11, 2025

Ngone51 Jan 13, 2025

xuanyuanking left a comment

Ngone51 commented Jan 15, 2025

LuciferYang commented Jan 15, 2025

Ngone51 commented Jan 15, 2025

LuciferYang commented Jan 15, 2025

Ngone51 commented Jan 15, 2025

Ngone51 commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

dongjoon-hyun commented Jan 15, 2025

dongjoon-hyun commented Jan 15, 2025

dongjoon-hyun Jan 15, 2025

Ngone51 Jan 15, 2025

Ngone51 Jan 15, 2025

dongjoon-hyun Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

dongjoon-hyun commented Jan 15, 2025

LuciferYang commented Jan 15, 2025 •

edited

Loading

HyukjinKwon commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

LuciferYang commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

cloud-fan Jan 15, 2025

LuciferYang Jan 15, 2025

Ngone51 Jan 15, 2025

cloud-fan Jan 15, 2025

Ngone51 Jan 15, 2025

cloud-fan Jan 15, 2025

Ngone51 Jan 15, 2025

Ngone51 commented Jan 15, 2025


		override def interruptible(): Boolean = TaskContext.synchronized(_interruptible)

		override def pendingInterrupt(threadToInterrupt: Option[Thread], reason: String): Unit = {

[SPARK-50768][CORE] Introduce TaskContext.createResourceUninterruptibly to avoid stream leak by task interruption #49413

[SPARK-50768][CORE] Introduce TaskContext.createResourceUninterruptibly to avoid stream leak by task interruption #49413

Conversation

Ngone51 commented Jan 8, 2025 • edited Loading

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xuanyuanking left a comment

Choose a reason for hiding this comment

Ngone51 commented Jan 15, 2025

LuciferYang commented Jan 15, 2025

Ngone51 commented Jan 15, 2025

LuciferYang commented Jan 15, 2025

Ngone51 commented Jan 15, 2025

Ngone51 commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

dongjoon-hyun commented Jan 15, 2025

dongjoon-hyun commented Jan 15, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HyukjinKwon commented Jan 15, 2025

dongjoon-hyun commented Jan 15, 2025

LuciferYang commented Jan 15, 2025 • edited Loading

HyukjinKwon commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

LuciferYang commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

HyukjinKwon commented Jan 15, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ngone51 commented Jan 15, 2025

Ngone51 commented Jan 8, 2025 •

edited

Loading

LuciferYang commented Jan 15, 2025 •

edited

Loading