This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 8f56b0ec83d6 [SPARK-58192][CORE][FOLLOWUP] Preserve exact task CPU
amounts in PySpark
8f56b0ec83d6 is described below
commit 8f56b0ec83d6e386db77bbcb68c2c6f348783a36
Author: Wenchen Fan <[email protected]>
AuthorDate: Fri Aug 7 00:37:59 2026 +0800
[SPARK-58192][CORE][FOLLOWUP] Preserve exact task CPU amounts in PySpark
### What changes were proposed in this pull request?
Follow-up to https://github.com/apache/spark/pull/57332.
This PR changes PySpark `TaskContext.cpuAmount()` to return
`decimal.Decimal` instead of
`float`, preserving the exact task CPU amount sent by the JVM. It also
updates the corresponding
worker message type and tests, clarifies the valid fractional CPU range
after rounding, and fixes
wording issues in comments added by the original change.
### Why are the changes needed?
Converting the exact decimal CPU amount to a Python `float` can lose
precision. Returning
`decimal.Decimal` keeps the Python API consistent with the exact
`BigDecimal` representation
used by the JVM.
### Does this PR introduce _any_ user-facing change?
Yes. In the unreleased fractional CPU API, PySpark
`TaskContext.cpuAmount()` now returns
`decimal.Decimal` instead of `float`.
### How was this patch tested?
Updated `pyspark.tests.test_taskcontext` to verify integer and
9-decimal-place fractional CPU
amounts as `decimal.Decimal` values.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #57810 from cloud-fan/SPARK-58192-followup.
Authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit da5f016c7e41d8fb5db4f3105e7879f8c6cddd5e)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../src/main/scala/org/apache/spark/executor/Executor.scala | 2 +-
.../org/apache/spark/resource/TaskResourceRequest.scala | 6 +++---
.../org/apache/spark/scheduler/TaskSchedulerImpl.scala | 4 ++--
python/pyspark/taskcontext.py | 13 +++++++------
python/pyspark/tests/test_taskcontext.py | 7 ++++---
python/pyspark/worker_message.py | 5 +++--
.../spark/sql/execution/BaseScriptTransformationExec.scala | 2 +-
7 files changed, 21 insertions(+), 18 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
index a67e700bc372..3056c3b38ad6 100644
--- a/core/src/main/scala/org/apache/spark/executor/Executor.scala
+++ b/core/src/main/scala/org/apache/spark/executor/Executor.scala
@@ -1611,7 +1611,7 @@ private[spark] object Executor extends Logging {
* minus executorRunTime -- would misreport the un-consumed share of the
core as scheduler
* delay. Relative comparisons for speculation are unaffected because all
tasks of a stage
* share the same cpu amount. Trailing zeros are stripped from the scale-9
cpus so the
- * product stays in BigDecimal's compact (long-backed) form instead of
inflating a long
+ * product stays in BigDecimal's compact (long-backed) form instead of
inflating a long
* duration into a BigInteger.
*/
private[spark] def cpuWeightedNanos(intervalNs: Long, cpus: BigDecimal):
Long = {
diff --git
a/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala
b/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala
index cdf1a1ec7bce..8341a723508d 100644
--- a/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala
+++ b/core/src/main/scala/org/apache/spark/resource/TaskResourceRequest.scala
@@ -32,9 +32,9 @@ import org.apache.spark.annotation.{Since, Stable}
* numbers, since a task's amount must map onto discrete
resource addresses -
* ie amount equals 0.5 translates into 2 tasks per resource
address. CPUs
* (resource name "cpus") are a plain quantity drawn from the
executor's core
- * pool rather than an addressable resource, so any amount of at
least 1e-9 is
- * valid, e.g. 1.5; the cpus amount is rounded to the nearest
1e-9, so precision
- * beyond 9 decimal places is not preserved.
+ * pool rather than an addressable resource, so any amount from
1e-9 through
+ * Int.MaxValue after rounding is valid, e.g. 1.5; the cpus
amount is rounded to
+ * the nearest 1e-9, so precision beyond 9 decimal places is not
preserved.
*/
@Stable
@Since("3.1.0")
diff --git
a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
index e5fdfc68701f..887545b5da1b 100644
--- a/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
+++ b/core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala
@@ -520,7 +520,7 @@ private[spark] class TaskSchedulerImpl(
availWorkerResources: ExecutorResourcesAmounts): Option[Map[String,
Map[String, Long]]] = {
val rpId = taskSet.taskSet.resourceProfileId
val taskSetProf = sc.resourceProfileManager.resourceProfileFromId(rpId)
- // check if the ResourceProfile has cpus first since that is common case.
Both values are in
+ // check if the ResourceProfile has cpus first since that is the common
case. Both values are in
// the internal exact BigDecimal representation, so this comparison is
exact regardless of
// whether spark.task.cpus is fractional (e.g. 0.2).
if (availCpus < taskCpus) return None
@@ -585,7 +585,7 @@ private[spark] class TaskSchedulerImpl(
val shuffledOffers = shuffleOffers(filteredOffers)
// Build a list of tasks to assign to each worker.
// Note the size estimate here might be off with different
ResourceProfiles but should be
- // close estimate. It is only a capacity hint, so cap it: with a tiny
fractional
+ // a close estimate. It is only a capacity hint, so cap it: with a tiny
fractional
// spark.task.cpus the exact slot count can be huge and would preallocate
a giant buffer.
val tasks = shuffledOffers.map { o =>
val sizeHint =
diff --git a/python/pyspark/taskcontext.py b/python/pyspark/taskcontext.py
index 796e3b5ef5f6..7e1dd2dc793d 100644
--- a/python/pyspark/taskcontext.py
+++ b/python/pyspark/taskcontext.py
@@ -14,6 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
+from decimal import Decimal
from typing import Any, ClassVar, Type, TypeVar, Dict, List, Optional, Union,
cast
from pyspark.util import local_connect_and_auth
@@ -128,7 +129,7 @@ class TaskContext:
_taskAttemptId: Optional[int] = None
_localProperties: Optional[Dict[str, str]] = None
_cpus: Optional[int] = None
- _cpuAmount: Optional[float] = None
+ _cpuAmount: Optional[Decimal] = None
_resources: Optional[Dict[str, "ResourceInformation"]] = None
def __new__(cls: Type["TaskContext"], **kwargs: Any) -> "TaskContext":
@@ -192,7 +193,7 @@ class TaskContext:
attemptNumber=json["attemptNumber"],
taskAttemptId=json["taskAttemptId"],
cpus=json["cpus"],
- cpuAmount=float(json["cpuAmount"]) if "cpuAmount" in json else
None,
+ cpuAmount=Decimal(json["cpuAmount"]) if "cpuAmount" in json else
None,
resources={
k: ResourceInformation(v["name"], v["addresses"])
for k, v in json["resources"].items()
@@ -276,7 +277,7 @@ class TaskContext:
"""
return cast(int, self._cpus)
- def cpuAmount(self) -> float:
+ def cpuAmount(self) -> Decimal:
"""
The exact amount of CPUs allocated to the task. This can be fractional
when
``spark.task.cpus`` or the task resource profile requests a fractional
amount
@@ -286,14 +287,14 @@ class TaskContext:
Returns
-------
- float
- the exact, possibly fractional, amount of CPUs.
+ decimal.Decimal
+ the exact, possibly fractional amount of CPUs.
See Also
--------
TaskContext.cpus
"""
- return cast(float, self._cpuAmount)
+ return cast(Decimal, self._cpuAmount)
def resources(self) -> Dict[str, "ResourceInformation"]:
"""
diff --git a/python/pyspark/tests/test_taskcontext.py
b/python/pyspark/tests/test_taskcontext.py
index 2f39ff53bc9e..8383f0052638 100644
--- a/python/pyspark/tests/test_taskcontext.py
+++ b/python/pyspark/tests/test_taskcontext.py
@@ -15,6 +15,7 @@
# limitations under the License.
#
import os
+from decimal import Decimal
import random
import socketserver
import stat
@@ -355,7 +356,7 @@ class TaskContextTestsWithResources(unittest.TestCase):
"""SPARK-58192: the exact cpu amount is available."""
rdd = self.sc.parallelize(range(10))
cpu_amount = rdd.map(lambda x:
TaskContext.get().cpuAmount()).take(1)[0]
- self.assertEqual(cpu_amount, 2.0)
+ self.assertEqual(cpu_amount, Decimal("2"))
def test_resources(self):
"""Test that multiple resources are all available (SPARK-54929)."""
@@ -379,7 +380,7 @@ class TaskContextTestsWithFractionalCpus(unittest.TestCase):
def setUp(self):
class_name = self.__class__.__name__
conf = SparkConf().set("spark.test.home", SPARK_HOME)
- conf = conf.set("spark.task.cpus", "0.5")
+ conf = conf.set("spark.task.cpus", "0.123456789")
self.sc = SparkContext("local-cluster[1,2,1024]", class_name,
conf=conf)
def test_fractional_cpu_amount(self):
@@ -388,7 +389,7 @@ class TaskContextTestsWithFractionalCpus(unittest.TestCase):
cpu_amount, cpus = rdd.map(
lambda x: (TaskContext.get().cpuAmount(), TaskContext.get().cpus())
).take(1)[0]
- self.assertEqual(cpu_amount, 0.5)
+ self.assertEqual(cpu_amount, Decimal("0.123456789"))
self.assertEqual(cpus, 1)
def test_omp_num_threads_follows_task_cpus(self):
diff --git a/python/pyspark/worker_message.py b/python/pyspark/worker_message.py
index 5be806882c97..b15d304f565d 100644
--- a/python/pyspark/worker_message.py
+++ b/python/pyspark/worker_message.py
@@ -16,6 +16,7 @@
#
import dataclasses
+from decimal import Decimal
import json
import sys
from typing import Optional, TypeAlias, Union, IO, Any
@@ -43,7 +44,7 @@ class TaskContextInfo:
attempt_number: int
task_attempt_id: int
cpus: int
- cpu_amount: float
+ cpu_amount: Decimal
resources: dict[str, ResourceInfo]
local_properties: dict[str, str]
@@ -59,7 +60,7 @@ class TaskContextInfo:
attempt_number=task_context_json["attemptNumber"],
task_attempt_id=task_context_json["taskAttemptId"],
cpus=task_context_json["cpus"],
- cpu_amount=float(task_context_json["cpuAmount"]),
+ cpu_amount=Decimal(task_context_json["cpuAmount"]),
resources={
k: cls.ResourceInfo(name=v["name"], addresses=v["addresses"])
for k, v in task_context_json["resources"].items()
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
index 2d8fec806f7f..a313e2c671be 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/BaseScriptTransformationExec.scala
@@ -85,7 +85,7 @@ trait BaseScriptTransformationExec extends UnaryExecNode {
SparkFiles.getRootDirectory()
builder.environment().put("PATH", path)
// if OMP_NUM_THREADS is not explicitly set, override it with the value of
"spark.task.cpus"
- // which may be fractional, so round up to an integer (at least 1) of
threads.
+ // which may be fractional, so round up to an integer number of threads
(at least 1).
if (System.getenv("OMP_NUM_THREADS") == null) {
val taskCpus = conf.getConfString("spark.task.cpus", "1.0").toDouble
builder.environment().put("OMP_NUM_THREADS", math.max(1,
taskCpus.ceil.toInt).toString)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
