(spark) branch branch-4.3 updated: [SPARK-59060][PYTHON] Add missing cpuAmount to the benchmark mock task context

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

viirya pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.3 by this push:
     new 2cee1398a63f [SPARK-59060][PYTHON] Add missing cpuAmount to the benchmark mock task context
2cee1398a63f is described below

commit 2cee1398a63fb4bd4411b49b0cecbf7ef31c445d
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Thu Aug 27 14:57:33 2026 -0700

    [SPARK-59060][PYTHON] Add missing cpuAmount to the benchmark mock task context
    
    ### What changes were proposed in this pull request?
    
    Adds the `cpuAmount` key to the mock task-context JSON that `MockProtocolWriter` writes in `python/benchmarks/bench_eval_type.py`.
    
    ### Why are the changes needed?
    
    Every benchmark in `bench_eval_type.py` currently fails at worker initialization. `TaskContextInfo.from_stream` reads `task_context_json["cpuAmount"]` without a default (`python/pyspark/worker_message.py`), so the worker raises `KeyError: 'cpuAmount'` before any UDF runs:
    
    ```
      File "python/pyspark/worker_message.py", line 63, in from_stream
        cpu_amount=Decimal(task_context_json["cpuAmount"]),
    KeyError: 'cpuAmount'
    ```
    
    The key became required in SPARK-58192, which added it on the JVM side and in `worker_message.py` but did not update this mock.
    
    The failure is easy to miss because `setup()` only builds the input buffer; only running `time_worker` / `peakmem_worker` reaches the worker. Nothing under `python/benchmarks/` is listed in `dev/sparktestsupport/modules.py`, so no CI job exercises these benchmarks.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No. Benchmark-only change.
    
    ### How was this patch tested?
    
    Ran every benchmark class in the file, driving the worker end to end:
    
    | | before | after |
    |---|---|---|
    | `*TimeBench` | 0 of 26 | 26 of 26 |
    | `*PeakmemBench` | 0 of 26 | 26 of 26 |
    
    Also checked out each release branch and ran the same benchmarks there:
    
    | branch | result |
    |---|---|
    | `master` | 0 of 26; 26 of 26 with this fix |
    | `branch-4.x` | same failure; the same one-line fix applies |
    | `branch-4.3` | 0 of 26; 26 of 26 with this fix |
    | `branch-4.2` | unaffected -- no `worker_message.py`, so the key is not required |
    
    So this is worth backporting to `branch-4.x` and `branch-4.3`.
    
    The value is a plain decimal string, matching `CpuAmount.toDisplayString` on the JVM side (`core/src/main/scala/org/apache/spark/api/python/PythonWorkerUtils.scala`), and is consistent with the existing `"cpus": 1` in the same mock.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Code (Claude Opus 5)
    
    Closes #58353 from viirya/SPARK-bench-cpuamount.
    
    Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
    Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
    (cherry picked from commit c5eebb341db59ad82dd013ff0067c8f742b0907b)
    Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
 python/benchmarks/bench_eval_type.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py
index 3228d915f44e..473de1d2924c 100644
--- a/python/benchmarks/bench_eval_type.py
+++ b/python/benchmarks/bench_eval_type.py
@@ -119,6 +119,9 @@ class MockProtocolWriter:
                     "attemptNumber": 0,
                     "taskAttemptId": 0,
                     "cpus": 1,
+                    # Plain decimal string, matching CpuAmount.toDisplayString on
+                    # the JVM side; TaskContextInfo.from_stream requires the key.
+                    "cpuAmount": "1",
                     "resources": {},
                     "localProperties": {},
                 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org