set env in Task

sunchao · Feb 23, 2024 · ac21273 · ac21273
1 parent a6502e0
commit ac21273
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/core/src/main/scala/org/apache/spark/executor/Executor.scala b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -589,6 +589,7 @@ private[spark] class Executor(
           taskDescription.serializedTask, Thread.currentThread.getContextClassLoader)
         task.localProperties = taskDescription.properties
         task.setTaskMemoryManager(taskMemoryManager)
+        task.setEnv(env)
 
         // If this task has been killed before we deserialized it, let's quit now. Otherwise,
         // continue executing the task.

diff --git a/core/src/main/scala/org/apache/spark/scheduler/Task.scala b/core/src/main/scala/org/apache/spark/scheduler/Task.scala
@@ -126,7 +126,7 @@ private[spark] abstract class Task[T](
 
     new CallerContext(
       "TASK",
-      SparkEnv.get.conf.get(APP_CALLER_CONTEXT),
+      env.conf.get(APP_CALLER_CONTEXT),
       appId,
       appAttemptId,
       jobId,
@@ -143,15 +143,14 @@ private[spark] abstract class Task[T](
       try {
         Utils.tryLogNonFatalError {
           // Release memory used by this thread for unrolling blocks
-          SparkEnv.get.blockManager.memoryStore.releaseUnrollMemoryForThisTask(MemoryMode.ON_HEAP)
-          SparkEnv.get.blockManager.memoryStore.releaseUnrollMemoryForThisTask(
-            MemoryMode.OFF_HEAP)
+          env.blockManager.memoryStore.releaseUnrollMemoryForThisTask(MemoryMode.ON_HEAP)
+          env.blockManager.memoryStore.releaseUnrollMemoryForThisTask(MemoryMode.OFF_HEAP)
           // Notify any tasks waiting for execution memory to be freed to wake up and try to
           // acquire memory again. This makes impossible the scenario where a task sleeps forever
           // because there are no other tasks left to notify it. Since this is safe to do but may
           // not be strictly necessary, we should revisit whether we can remove this in the
           // future.
-          val memoryManager = SparkEnv.get.memoryManager
+          val memoryManager = env.memoryManager
           memoryManager.synchronized { memoryManager.notifyAll() }
         }
       } finally {
@@ -164,11 +163,16 @@ private[spark] abstract class Task[T](
   }
 
   private var taskMemoryManager: TaskMemoryManager = _
+  private var env: SparkEnv = _
 
   def setTaskMemoryManager(taskMemoryManager: TaskMemoryManager): Unit = {
     this.taskMemoryManager = taskMemoryManager
   }
 
+  def setEnv(env: SparkEnv): Unit = {
+    this.env = env
+  }
+
   def runTask(context: TaskContext): T
 
   def preferredLocations: Seq[TaskLocation] = Nil