python/pyspark/tests/test_taskcontext.py - master - githubmirror/spark

python/pyspark/tests/test_taskcontext.py
417 строк17 KB
Wenchen Fan
[SPARK-58192][CORE][FOLLOWUP] Preserve exact task CPU amounts in PySpark
06 авг 2026, 19:37
06 авг 2026, 19:37da5f016
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import os
from decimal import Decimal
import random
import socketserver
import stat
import sys
import tempfile
import time
import unittest

from pyspark import SparkConf, SparkContext, TaskContext, BarrierTaskContext
from pyspark.testing.sqlutils import SPARK_HOME
from pyspark.testing.utils import PySparkTestCase, eventually


class TaskContextTests(PySparkTestCase):
    def setUp(self):
        self._old_sys_path = list(sys.path)
        class_name = self.__class__.__name__
        # Allow retries even though they are normally disabled in local mode
        self.sc = SparkContext("local[4, 2]", class_name)

    def test_stage_id(self):
        """Test the stage ids are available and incrementing as expected."""
        rdd = self.sc.parallelize(range(10))
        stage1 = rdd.map(lambda x: TaskContext.get().stageId()).take(1)[0]
        stage2 = rdd.map(lambda x: TaskContext.get().stageId()).take(1)[0]
        # Test using the constructor directly rather than the get()
        stage3 = rdd.map(lambda x: TaskContext().stageId()).take(1)[0]
        self.assertEqual(stage1 + 1, stage2)
        self.assertEqual(stage1 + 2, stage3)
        self.assertEqual(stage2 + 1, stage3)

    def test_resources(self):
        """Test the resources are empty by default."""
        rdd = self.sc.parallelize(range(10))
        resources1 = rdd.map(lambda x: TaskContext.get().resources()).take(1)[0]
        # Test using the constructor directly rather than the get()
        resources2 = rdd.map(lambda x: TaskContext().resources()).take(1)[0]
        self.assertEqual(len(resources1), 0)
        self.assertEqual(len(resources2), 0)

    def test_partition_id(self):
        """Test the partition id."""
        rdd1 = self.sc.parallelize(range(10), 1)
        rdd2 = self.sc.parallelize(range(10), 2)
        pids1 = rdd1.map(lambda x: TaskContext.get().partitionId()).collect()
        pids2 = rdd2.map(lambda x: TaskContext.get().partitionId()).collect()
        self.assertEqual(0, pids1[0])
        self.assertEqual(0, pids1[9])
        self.assertEqual(0, pids2[0])
        self.assertEqual(1, pids2[9])

    def test_attempt_number(self):
        """Verify the attempt numbers are correctly reported."""
        rdd = self.sc.parallelize(range(10))
        # Verify a simple job with no failures
        attempt_numbers = rdd.map(lambda x: TaskContext.get().attemptNumber()).collect()
        map(lambda attempt: self.assertEqual(0, attempt), attempt_numbers)

        def fail_on_first(x):
            """Fail on the first attempt so we get a positive attempt number"""
            tc = TaskContext.get()
            attempt_number = tc.attemptNumber()
            partition_id = tc.partitionId()
            attempt_id = tc.taskAttemptId()
            if attempt_number == 0 and partition_id == 0:
                raise RuntimeError("Failing on first attempt")
            else:
                return [x, partition_id, attempt_number, attempt_id]

        result = rdd.map(fail_on_first).collect()
        # We should re-submit the first partition to it but other partitions should be attempt 0
        self.assertEqual([0, 0, 1], result[0][0:3])
        self.assertEqual([9, 3, 0], result[9][0:3])
        first_partition = filter(lambda x: x[1] == 0, result)
        map(lambda x: self.assertEqual(1, x[2]), first_partition)
        other_partitions = filter(lambda x: x[1] != 0, result)
        map(lambda x: self.assertEqual(0, x[2]), other_partitions)
        # The task attempt id should be different
        self.assertTrue(result[0][3] != result[9][3])

    def test_tc_on_driver(self):
        """Verify that getting the TaskContext on the driver returns None."""
        tc = TaskContext.get()
        self.assertTrue(tc is None)

    def test_get_local_property(self):
        """Verify that local properties set on the driver are available in TaskContext."""
        key = "testkey"
        value = "testvalue"
        self.sc.setLocalProperty(key, value)
        try:
            rdd = self.sc.parallelize(range(1), 1)
            prop1 = rdd.map(lambda _: TaskContext.get().getLocalProperty(key)).collect()[0]
            self.assertEqual(prop1, value)
            prop2 = rdd.map(lambda _: TaskContext.get().getLocalProperty("otherkey")).collect()[0]
            self.assertTrue(prop2 is None)
        finally:
            self.sc.setLocalProperty(key, None)

    def test_barrier(self):
        """
        Verify that BarrierTaskContext.barrier() performs global sync among all barrier tasks
        within a stage.
        """
        rdd = self.sc.parallelize(range(10), 4)

        def f(iterator):
            yield sum(iterator)

        def context_barrier(x):
            tc = BarrierTaskContext.get()
            time.sleep(random.randint(1, 5) * 2)
            tc.barrier()
            return time.time()

        times = rdd.barrier().mapPartitions(f).map(context_barrier).collect()
        self.assertTrue(max(times) - min(times) < 2)

    def test_all_gather(self):
        """
        Verify that BarrierTaskContext.allGather() performs global sync among all barrier tasks
        within a stage and passes messages properly.
        """
        rdd = self.sc.parallelize(range(10), 4)

        def f(iterator):
            yield sum(iterator)

        def context_barrier(x):
            tc = BarrierTaskContext.get()
            time.sleep(random.randint(1, 10))
            out = tc.allGather(str(tc.partitionId()))
            pids = [int(e) for e in out]
            return pids

        pids = rdd.barrier().mapPartitions(f).map(context_barrier).collect()[0]
        self.assertEqual(pids, [0, 1, 2, 3])

    def test_barrier_infos(self):
        """
        Verify that BarrierTaskContext.getTaskInfos() returns a list of all task infos in the
        barrier stage.
        """
        rdd = self.sc.parallelize(range(10), 4)

        def f(iterator):
            yield sum(iterator)

        taskInfos = (
            rdd.barrier()
            .mapPartitions(f)
            .map(lambda x: BarrierTaskContext.get().getTaskInfos())
            .collect()
        )
        self.assertEqual(len(taskInfos), 4)
        self.assertEqual(len(taskInfos[0]), 4)

    def test_context_get(self):
        """
        Verify that TaskContext.get() works both in or not in a barrier stage.
        """
        rdd = self.sc.parallelize(range(10), 4)

        def f(iterator):
            taskContext = TaskContext.get()
            if isinstance(taskContext, BarrierTaskContext):
                yield taskContext.partitionId() + 1
            elif isinstance(taskContext, TaskContext):
                yield taskContext.partitionId() + 2
            else:
                yield -1

        # for normal stage
        result1 = rdd.mapPartitions(f).collect()
        self.assertTrue(result1 == [2, 3, 4, 5])
        # for barrier stage
        result2 = rdd.barrier().mapPartitions(f).collect()
        self.assertTrue(result2 == [1, 2, 3, 4])

    def test_barrier_context_get(self):
        """
        Verify that BarrierTaskContext.get() should only works in a barrier stage.
        """
        rdd = self.sc.parallelize(range(10), 4)

        def f(iterator):
            try:
                taskContext = BarrierTaskContext.get()
            except Exception:
                yield -1
            else:
                yield taskContext.partitionId()

        # for normal stage
        result1 = rdd.mapPartitions(f).collect()
        self.assertTrue(result1 == [-1, -1, -1, -1])
        # for barrier stage
        result2 = rdd.barrier().mapPartitions(f).collect()
        self.assertTrue(result2 == [0, 1, 2, 3])


@unittest.skipUnless(
    hasattr(socketserver, "UnixStreamServer"),
    "Unix Domain Socket is not supported on this platform.",
)
class TaskContextUDSTests(TaskContextTests):
    def setUp(self):
        self._old_sys_path = list(sys.path)
        class_name = self.__class__.__name__
        # Enable Unix Domain Socket for the test
        conf = SparkConf().set("spark.python.unix.domain.socket.enabled", "true")
        # Allow retries even though they are normally disabled in local mode
        self.sc = SparkContext("local[4, 2]", class_name, conf=conf)


class TaskContextTestsWithWorkerReuse(unittest.TestCase):
    def setUp(self):
        class_name = self.__class__.__name__
        conf = SparkConf().set("spark.python.worker.reuse", "true")
        self.sc = SparkContext("local[2]", class_name, conf=conf)

    def test_barrier_with_python_worker_reuse(self):
        """
        Regression test for SPARK-25921: verify that BarrierTaskContext.barrier() with
        reused python worker.
        """
        # start a normal job first to start all workers and get all worker pids
        worker_pids = self.sc.parallelize(range(2), 2).map(lambda x: os.getpid()).collect()
        # the worker will reuse in this barrier job
        rdd = self.sc.parallelize(range(10), 2)

        def f(iterator):
            yield sum(iterator)

        def context_barrier(x):
            tc = BarrierTaskContext.get()
            time.sleep(random.randint(1, 5) * 2)
            tc.barrier()
            return (time.time(), os.getpid())

        result = rdd.barrier().mapPartitions(f).map(context_barrier).collect()
        times = list(map(lambda x: x[0], result))
        pids = list(map(lambda x: x[1], result))
        # check both barrier and worker reuse effect
        self.assertTrue(max(times) - min(times) < 2)
        for pid in pids:
            self.assertTrue(pid in worker_pids)

    def check_task_context_correct_with_python_worker_reuse(self):
        """Verify the task context correct when reused python worker"""
        # start a normal job first to start all workers and get all worker pids
        worker_pids = self.sc.parallelize(range(2), 2).map(lambda x: os.getpid()).collect()
        # the worker will reuse in this barrier job
        rdd = self.sc.parallelize(range(10), 2)

        def context(iterator):
            tp = TaskContext.get().partitionId()
            try:
                bp = BarrierTaskContext.get().partitionId()
            except Exception:
                bp = -1

            yield (tp, bp, os.getpid())

        # normal stage after normal stage
        normal_result = rdd.mapPartitions(context).collect()
        tps, bps, pids = zip(*normal_result)
        self.assertTrue(tps == (0, 1))
        self.assertTrue(bps == (-1, -1))
        for pid in pids:
            self.assertTrue(pid in worker_pids)
        # barrier stage after normal stage
        barrier_result = rdd.barrier().mapPartitions(context).collect()
        tps, bps, pids = zip(*barrier_result)
        self.assertTrue(tps == (0, 1))
        self.assertTrue(bps == (0, 1))
        for pid in pids:
            self.assertTrue(pid in worker_pids)
        # normal stage after barrier stage
        normal_result2 = rdd.mapPartitions(context).collect()
        tps, bps, pids = zip(*normal_result2)
        self.assertTrue(tps == (0, 1))
        self.assertTrue(bps == (-1, -1))
        for pid in pids:
            self.assertTrue(pid in worker_pids)
        return True

    @eventually(catch_assertions=True)
    def test_task_context_correct_with_python_worker_reuse(self):
        # Retrying the check as the PIDs from Python workers might be different even
        # when reusing Python workers is enabled if a Python worker is dead for some reasons
        # (e.g., socket connection failure) and new Python worker is created.
        self.check_task_context_correct_with_python_worker_reuse()

    def tearDown(self):
        self.sc.stop()


class TaskContextTestsWithResources(unittest.TestCase):
    def setUp(self):
        class_name = self.__class__.__name__
        # Create discovery script for GPU
        self.gpuTempFile = tempfile.NamedTemporaryFile(delete=False)
        self.gpuTempFile.write(b'echo {\\"name\\": \\"gpu\\", \\"addresses\\": [\\"0\\"]}')
        self.gpuTempFile.close()
        os.chmod(
            self.gpuTempFile.name,
            stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IROTH | stat.S_IXOTH,
        )
        # Create discovery script for FPGA (SPARK-54929)
        self.fpgaTempFile = tempfile.NamedTemporaryFile(delete=False)
        self.fpgaTempFile.write(b'echo {\\"name\\": \\"fpga\\", \\"addresses\\": [\\"0\\"]}')
        self.fpgaTempFile.close()
        os.chmod(
            self.fpgaTempFile.name,
            stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IROTH | stat.S_IXOTH,
        )
        conf = SparkConf().set("spark.test.home", SPARK_HOME)
        conf = conf.set("spark.worker.resource.gpu.discoveryScript", self.gpuTempFile.name)
        conf = conf.set("spark.worker.resource.gpu.amount", 1)
        conf = conf.set("spark.task.cpus", 2)
        conf = conf.set("spark.task.resource.gpu.amount", "1")
        conf = conf.set("spark.executor.resource.gpu.amount", "1")
        # Configure FPGA resource (SPARK-54929)
        conf = conf.set("spark.worker.resource.fpga.discoveryScript", self.fpgaTempFile.name)
        conf = conf.set("spark.worker.resource.fpga.amount", 1)
        conf = conf.set("spark.task.resource.fpga.amount", "1")
        conf = conf.set("spark.executor.resource.fpga.amount", "1")
        self.sc = SparkContext("local-cluster[2,2,1024]", class_name, conf=conf)

    def test_cpus(self):
        """Test the cpus are available."""
        rdd = self.sc.parallelize(range(10))
        cpus = rdd.map(lambda x: TaskContext.get().cpus()).take(1)[0]
        self.assertEqual(cpus, 2)

    def test_cpu_amount(self):
        """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, Decimal("2"))

    def test_resources(self):
        """Test that multiple resources are all available (SPARK-54929)."""
        rdd = self.sc.parallelize(range(10))
        resources = rdd.map(lambda x: TaskContext.get().resources()).take(1)[0]
        self.assertEqual(len(resources), 2)
        self.assertTrue("gpu" in resources)
        self.assertEqual(resources["gpu"].name, "gpu")
        self.assertEqual(resources["gpu"].addresses, ["0"])
        self.assertTrue("fpga" in resources)
        self.assertEqual(resources["fpga"].name, "fpga")
        self.assertEqual(resources["fpga"].addresses, ["0"])

    def tearDown(self):
        os.unlink(self.gpuTempFile.name)
        os.unlink(self.fpgaTempFile.name)
        self.sc.stop()


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.123456789")
        self.sc = SparkContext("local-cluster[1,2,1024]", class_name, conf=conf)

    def test_fractional_cpu_amount(self):
        """SPARK-58192: cpuAmount() exposes the exact fractional amount; cpus() rounds up."""
        rdd = self.sc.parallelize(range(4), 2)
        cpu_amount, cpus = rdd.map(
            lambda x: (TaskContext.get().cpuAmount(), TaskContext.get().cpus())
        ).take(1)[0]
        self.assertEqual(cpu_amount, Decimal("0.123456789"))
        self.assertEqual(cpus, 1)

    def test_omp_num_threads_follows_task_cpus(self):
        """SPARK-58192: OMP_NUM_THREADS defaults to the ceiling of the task's own cpu
        amount, honoring stage-level resource profiles over the global spark.task.cpus."""
        from pyspark.resource import ResourceProfileBuilder, TaskResourceRequests

        rdd = self.sc.parallelize(range(2), 1)
        omp = rdd.map(lambda _: os.environ["OMP_NUM_THREADS"]).take(1)[0]
        self.assertEqual(omp, "1")

        treqs = TaskResourceRequests().cpus(2)
        rp = ResourceProfileBuilder().require(treqs).build
        rdd = self.sc.parallelize(range(2), 1).withResources(rp)
        omp = rdd.map(lambda _: os.environ["OMP_NUM_THREADS"]).take(1)[0]
        self.assertEqual(omp, "2")

    def tearDown(self):
        self.sc.stop()


if __name__ == "__main__":
    from pyspark.testing import main

    main()