python/pyspark/sql/tests/test_job_cancellation.py - master - githubmirror/spark

python/pyspark/sql/tests/test_job_cancellation.py
231 строка8 KB
Jubin Soni
[SPARK-58358][SQL] Add tag validation to removeTag
07 авг 2026, 17:18
07 авг 2026, 17:188c22eff
#
# 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 unittest
import threading
import time

from pyspark import InheritableThread, inheritable_thread_target
from pyspark.errors import IllegalArgumentException, PySparkValueError
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import eventually


class JobCancellationTestsMixin:
    def test_tags(self):
        self.spark.clearTags()
        self.spark.addTag("a")
        self.assertEqual(self.spark.getTags(), {"a"})
        self.spark.addTag("b")
        self.spark.removeTag("a")
        self.assertEqual(self.spark.getTags(), {"b"})
        self.spark.addTag("c")
        self.spark.clearTags()
        self.assertEqual(self.spark.getTags(), set())
        self.spark.clearTags()

    def test_invalid_tags(self):
        # A tag cannot be an empty string or contain the ',' separator (documented on
        # SparkSession.addTag / removeTag). Both the classic and Spark Connect paths reject
        # such tags. The two paths raise different types -- a JVM-backed
        # IllegalArgumentException in Classic and a PySparkValueError in Connect -- so accept
        # either while still rejecting unrelated session/transport failures.
        invalid_tag_errors = (IllegalArgumentException, PySparkValueError)
        self.spark.clearTags()
        for invalid_tag in ["", "a,b", ","]:
            with self.assertRaises(invalid_tag_errors):
                self.spark.addTag(invalid_tag)
            with self.assertRaises(invalid_tag_errors):
                self.spark.removeTag(invalid_tag)
        # A rejected tag must not have been recorded.
        self.assertEqual(self.spark.getTags(), set())

        # Removing a valid but absent tag is a no-op in both modes (does not raise).
        self.spark.removeTag("absent")
        self.assertEqual(self.spark.getTags(), set())
        self.spark.clearTags()

    def test_tags_multithread(self):
        output1 = None
        output2 = None

        def tag1():
            nonlocal output1

            self.spark.addTag("tag1")
            output1 = self.spark.getTags()

        def tag2():
            nonlocal output2

            self.spark.addTag("tag2")
            output2 = self.spark.getTags()

        t1 = threading.Thread(target=tag1)
        t1.start()
        t1.join()
        t2 = threading.Thread(target=tag2)
        t2.start()
        t2.join()

        self.assertIsNotNone(output1)
        self.assertEqual(output1, {"tag1"})
        self.assertIsNotNone(output2)
        self.assertEqual(output2, {"tag2"})

    def check_job_cancellation(
        self, setter, canceller, thread_ids, thread_ids_to_cancel, thread_ids_to_run
    ):
        job_id_a = "job_ids_to_cancel"
        job_id_b = "job_ids_to_run"
        threads = []

        # A list which records whether job is cancelled.
        # The index of the array is the thread index which job run in.
        is_job_cancelled = [False for _ in thread_ids]

        def run_job(job_id, index, timeout=10):
            """
            Executes a job with the group ``job_group``. Each job waits for 3 seconds
            and then exits.
            """
            try:
                setter(job_id)

                from pyspark.sql.functions import udf
                from pyspark.sql.types import IntegerType

                def func(t: int) -> int:
                    time.sleep(t)
                    return t

                u = udf(func, IntegerType())

                self.spark.createDataFrame([[timeout]], ["a"]).repartition(1).select(
                    u("a").alias("b")
                ).collect()
                is_job_cancelled[index] = False
            except Exception:
                # Assume that exception means job cancellation.
                is_job_cancelled[index] = True

        # Test if job succeeded when not cancelled.
        run_job(job_id_a, 0, timeout=0)
        self.assertFalse(is_job_cancelled[0])
        self.spark.clearTags()

        # Run jobs
        for i in thread_ids_to_cancel:
            t = threading.Thread(target=run_job, args=(job_id_a, i))
            t.start()
            threads.append(t)

        for i in thread_ids_to_run:
            t = threading.Thread(target=run_job, args=(job_id_b, i))
            t.start()
            threads.append(t)

        @eventually(timeout=10, catch_assertions=True)
        def cancel_until_cancelled(cancel_ids):
            canceller(job_id_a)
            for cancel_id in cancel_ids:
                self.assertTrue(is_job_cancelled[cancel_id])
            return True

        cancel_until_cancelled(thread_ids_to_cancel)

        # Wait until all threads launching jobs are finished.
        for t in threads:
            t.join()

        for i in thread_ids_to_cancel:
            self.assertTrue(
                is_job_cancelled[i], "Thread {i}: Job in group A was not cancelled.".format(i=i)
            )

        for i in thread_ids_to_run:
            self.assertFalse(
                is_job_cancelled[i], "Thread {i}: Job in group B did not succeeded.".format(i=i)
            )

    def test_inheritable_tags(self):
        self.check_inheritable_tags(
            create_thread=lambda target, session: InheritableThread(target, session=session)
        )
        self.check_inheritable_tags(
            create_thread=lambda target, session: threading.Thread(
                target=inheritable_thread_target(session)(target)
            )
        )

    def check_inheritable_tags(self, create_thread):
        spark = self.spark
        spark.addTag("a")
        first = set()
        second = set()

        def get_inner_local_prop():
            spark.addTag("c")
            second.update(spark.getTags())

        def get_outer_local_prop():
            spark.addTag("b")
            first.update(spark.getTags())
            t2 = create_thread(target=get_inner_local_prop, session=spark)
            t2.start()
            t2.join()

        t1 = create_thread(target=get_outer_local_prop, session=spark)
        t1.start()
        t1.join()

        self.assertEqual(spark.getTags(), {"a"})
        self.assertEqual(first, {"a", "b"})
        self.assertEqual(second, {"a", "b", "c"})

    @unittest.skip("SPARK-52054: Should enable this test")
    def test_interrupt_tag(self):
        thread_ids = range(4)
        self.check_job_cancellation(
            lambda job_group: self.spark.addTag(job_group),
            lambda job_group: self.spark.interruptTag(job_group),
            thread_ids,
            [i for i in thread_ids if i % 2 == 0],
            [i for i in thread_ids if i % 2 != 0],
        )
        self.spark.clearTags()

    def test_interrupt_all(self):
        thread_ids = range(4)
        self.check_job_cancellation(
            lambda job_group: None,
            lambda job_group: self.spark.interruptAll(),
            thread_ids,
            thread_ids,
            [],
        )
        self.spark.clearTags()


class JobCancellationTests(JobCancellationTestsMixin, ReusedSQLTestCase):
    pass


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

    main()