snowpark-connect 0.24.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/tests/test_serializers.py

@@ -1,257 +0,0 @@
-#
-# 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 math
-import sys
-import unittest
-
-from pyspark import serializers
-from pyspark.serializers import (
-    CloudPickleSerializer,
-    CompressedSerializer,
-    AutoBatchedSerializer,
-    BatchedSerializer,
-    AutoSerializer,
-    NoOpSerializer,
-    PairDeserializer,
-    FlattenedValuesSerializer,
-    CartesianDeserializer,
-    CPickleSerializer,
-    UTF8Deserializer,
-    MarshalSerializer,
-)
-from pyspark.testing.utils import (
-    PySparkTestCase,
-    read_int,
-    write_int,
-    ByteArrayOutput,
-    have_numpy,
-    have_scipy,
-)
-
-
-class SerializationTestCase(unittest.TestCase):
-    def test_namedtuple(self):
-        from collections import namedtuple
-        from pyspark.cloudpickle import dumps, loads
-
-        P = namedtuple("P", "x y")
-        p1 = P(1, 3)
-        p2 = loads(dumps(p1, 2))
-        self.assertEqual(p1, p2)
-
-        P2 = loads(dumps(P))
-        p3 = P2(1, 3)
-        self.assertEqual(p1, p3)
-
-    def test_itemgetter(self):
-        from operator import itemgetter
-
-        ser = CloudPickleSerializer()
-        d = range(10)
-        getter = itemgetter(1)
-        getter2 = ser.loads(ser.dumps(getter))
-        self.assertEqual(getter(d), getter2(d))
-
-        getter = itemgetter(0, 3)
-        getter2 = ser.loads(ser.dumps(getter))
-        self.assertEqual(getter(d), getter2(d))
-
-    def test_function_module_name(self):
-        ser = CloudPickleSerializer()
-
-        def func(x):
-            return x
-
-        func2 = ser.loads(ser.dumps(func))
-        self.assertEqual(func.__module__, func2.__module__)
-
-    def test_attrgetter(self):
-        from operator import attrgetter
-
-        ser = CloudPickleSerializer()
-
-        class C:
-            def __getattr__(self, item):
-                return item
-
-        d = C()
-        getter = attrgetter("a")
-        getter2 = ser.loads(ser.dumps(getter))
-        self.assertEqual(getter(d), getter2(d))
-        getter = attrgetter("a", "b")
-        getter2 = ser.loads(ser.dumps(getter))
-        self.assertEqual(getter(d), getter2(d))
-
-        d.e = C()
-        getter = attrgetter("e.a")
-        getter2 = ser.loads(ser.dumps(getter))
-        self.assertEqual(getter(d), getter2(d))
-        getter = attrgetter("e.a", "e.b")
-        getter2 = ser.loads(ser.dumps(getter))
-        self.assertEqual(getter(d), getter2(d))
-
-    # Regression test for SPARK-3415
-    def test_pickling_file_handles(self):
-        # to be corrected with SPARK-11160
-        try:
-            import xmlrunner  # noqa: F401
-        except ImportError:
-            ser = CloudPickleSerializer()
-            out1 = sys.stderr
-            out2 = ser.loads(ser.dumps(out1))
-            self.assertEqual(out1, out2)
-
-    def test_func_globals(self):
-        class Unpicklable:
-            def __reduce__(self):
-                raise RuntimeError("not picklable")
-
-        global exit
-        exit = Unpicklable()
-
-        ser = CloudPickleSerializer()
-        self.assertRaises(Exception, lambda: ser.dumps(exit))
-
-        def foo():
-            sys.exit(0)
-
-        self.assertTrue("exit" in foo.__code__.co_names)
-        ser.dumps(foo)
-
-    def test_compressed_serializer(self):
-        ser = CompressedSerializer(CPickleSerializer())
-        from io import BytesIO as StringIO
-
-        io = StringIO()
-        ser.dump_stream(["abc", "123", range(5)], io)
-        io.seek(0)
-        self.assertEqual(["abc", "123", range(5)], list(ser.load_stream(io)))
-        ser.dump_stream(range(1000), io)
-        io.seek(0)
-        self.assertEqual(["abc", "123", range(5)] + list(range(1000)), list(ser.load_stream(io)))
-        io.close()
-
-    def test_hash_serializer(self):
-        hash(NoOpSerializer())
-        hash(UTF8Deserializer())
-        hash(CPickleSerializer())
-        hash(MarshalSerializer())
-        hash(AutoSerializer())
-        hash(BatchedSerializer(CPickleSerializer()))
-        hash(AutoBatchedSerializer(MarshalSerializer()))
-        hash(PairDeserializer(NoOpSerializer(), UTF8Deserializer()))
-        hash(CartesianDeserializer(NoOpSerializer(), UTF8Deserializer()))
-        hash(CompressedSerializer(CPickleSerializer()))
-        hash(FlattenedValuesSerializer(CPickleSerializer()))
-
-
-@unittest.skipIf(not have_scipy, "SciPy not installed")
-class SciPyTests(PySparkTestCase):
-
-    """General PySpark tests that depend on scipy"""
-
-    def test_serialize(self):
-        from scipy.special import gammaln
-
-        x = range(1, 5)
-        expected = list(map(gammaln, x))
-        observed = self.sc.parallelize(x).map(gammaln).collect()
-        self.assertEqual(expected, observed)
-
-
-@unittest.skipIf(not have_numpy, "NumPy not installed")
-class NumPyTests(PySparkTestCase):
-
-    """General PySpark tests that depend on numpy"""
-
-    def test_statcounter_array(self):
-        import numpy as np
-
-        x = self.sc.parallelize([np.array([1.0, 1.0]), np.array([2.0, 2.0]), np.array([3.0, 3.0])])
-        s = x.stats()
-        self.assertSequenceEqual([2.0, 2.0], s.mean().tolist())
-        self.assertSequenceEqual([1.0, 1.0], s.min().tolist())
-        self.assertSequenceEqual([3.0, 3.0], s.max().tolist())
-        self.assertSequenceEqual([1.0, 1.0], s.sampleStdev().tolist())
-
-        stats_dict = s.asDict()
-        self.assertEqual(3, stats_dict["count"])
-        self.assertSequenceEqual([2.0, 2.0], stats_dict["mean"].tolist())
-        self.assertSequenceEqual([1.0, 1.0], stats_dict["min"].tolist())
-        self.assertSequenceEqual([3.0, 3.0], stats_dict["max"].tolist())
-        self.assertSequenceEqual([6.0, 6.0], stats_dict["sum"].tolist())
-        self.assertSequenceEqual([1.0, 1.0], stats_dict["stdev"].tolist())
-        self.assertSequenceEqual([1.0, 1.0], stats_dict["variance"].tolist())
-
-        stats_sample_dict = s.asDict(sample=True)
-        self.assertEqual(3, stats_dict["count"])
-        self.assertSequenceEqual([2.0, 2.0], stats_sample_dict["mean"].tolist())
-        self.assertSequenceEqual([1.0, 1.0], stats_sample_dict["min"].tolist())
-        self.assertSequenceEqual([3.0, 3.0], stats_sample_dict["max"].tolist())
-        self.assertSequenceEqual([6.0, 6.0], stats_sample_dict["sum"].tolist())
-        self.assertSequenceEqual(
-            [0.816496580927726, 0.816496580927726], stats_sample_dict["stdev"].tolist()
-        )
-        self.assertSequenceEqual(
-            [0.6666666666666666, 0.6666666666666666], stats_sample_dict["variance"].tolist()
-        )
-
-
-class SerializersTest(unittest.TestCase):
-    def test_chunked_stream(self):
-        original_bytes = bytearray(range(100))
-        for data_length in [1, 10, 100]:
-            for buffer_length in [1, 2, 3, 5, 20, 99, 100, 101, 500]:
-                dest = ByteArrayOutput()
-                stream_out = serializers.ChunkedStream(dest, buffer_length)
-                stream_out.write(original_bytes[:data_length])
-                stream_out.close()
-                num_chunks = int(math.ceil(float(data_length) / buffer_length))
-                # length for each chunk, and a final -1 at the very end
-                exp_size = (num_chunks + 1) * 4 + data_length
-                self.assertEqual(len(dest.buffer), exp_size)
-                dest_pos = 0
-                data_pos = 0
-                for chunk_idx in range(num_chunks):
-                    chunk_length = read_int(dest.buffer[dest_pos : (dest_pos + 4)])
-                    if chunk_idx == num_chunks - 1:
-                        exp_length = data_length % buffer_length
-                        if exp_length == 0:
-                            exp_length = buffer_length
-                    else:
-                        exp_length = buffer_length
-                    self.assertEqual(chunk_length, exp_length)
-                    dest_pos += 4
-                    dest_chunk = dest.buffer[dest_pos : dest_pos + chunk_length]
-                    orig_chunk = original_bytes[data_pos : data_pos + chunk_length]
-                    self.assertEqual(dest_chunk, orig_chunk)
-                    dest_pos += chunk_length
-                    data_pos += chunk_length
-                # ends with a -1
-                self.assertEqual(dest.buffer[-4:], write_int(-1))
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_serializers import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/tests/test_shuffle.py

@@ -1,267 +0,0 @@
-#
-# 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 random
-import unittest
-import tempfile
-import os
-
-from py4j.protocol import Py4JJavaError
-
-from pyspark import shuffle, CPickleSerializer, SparkConf, SparkContext
-from pyspark.shuffle import (
-    Aggregator,
-    ExternalMerger,
-    ExternalSorter,
-    SimpleAggregator,
-    Merger,
-    ExternalGroupBy,
-)
-
-
-class MergerTests(unittest.TestCase):
-    def setUp(self):
-        self.N = 1 << 12
-        self.lst = [i for i in range(self.N)]
-        self.data = list(zip(self.lst, self.lst))
-        self.agg = Aggregator(
-            lambda x: [x], lambda x, y: x.append(y) or x, lambda x, y: x.extend(y) or x
-        )
-
-    def test_small_dataset(self):
-        m = ExternalMerger(self.agg, 1000)
-        m.mergeValues(self.data)
-        self.assertEqual(m.spills, 0)
-        self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N)))
-
-        m = ExternalMerger(self.agg, 1000)
-        m.mergeCombiners(map(lambda x_y1: (x_y1[0], [x_y1[1]]), self.data))
-        self.assertEqual(m.spills, 0)
-        self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N)))
-
-    def test_medium_dataset(self):
-        m = ExternalMerger(self.agg, 20)
-        m.mergeValues(self.data)
-        self.assertTrue(m.spills >= 1)
-        self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N)))
-
-        m = ExternalMerger(self.agg, 10)
-        m.mergeCombiners(map(lambda x_y2: (x_y2[0], [x_y2[1]]), self.data * 3))
-        self.assertTrue(m.spills >= 1)
-        self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N)) * 3)
-
-    def test_shuffle_data_with_multiple_locations(self):
-        # SPARK-39179: Test shuffle of data with multiple location also check
-        # shuffle locations get randomized
-
-        with tempfile.TemporaryDirectory() as tempdir1, tempfile.TemporaryDirectory() as tempdir2:
-            original = os.environ.get("SPARK_LOCAL_DIRS", None)
-            os.environ["SPARK_LOCAL_DIRS"] = tempdir1 + "," + tempdir2
-            try:
-                index_of_tempdir1 = [False, False]
-                for idx in range(10):
-                    m = ExternalMerger(self.agg, 20)
-                    if m.localdirs[0].startswith(tempdir1):
-                        index_of_tempdir1[0] = True
-                    elif m.localdirs[1].startswith(tempdir1):
-                        index_of_tempdir1[1] = True
-                    m.mergeValues(self.data)
-                    self.assertTrue(m.spills >= 1)
-                    self.assertEqual(sum(sum(v) for k, v in m.items()), sum(range(self.N)))
-                self.assertTrue(
-                    index_of_tempdir1[0] and (index_of_tempdir1[0] == index_of_tempdir1[1])
-                )
-            finally:
-                if original is not None:
-                    os.environ["SPARK_LOCAL_DIRS"] = original
-                else:
-                    del os.environ["SPARK_LOCAL_DIRS"]
-
-    def test_simple_aggregator_with_medium_dataset(self):
-        # SPARK-39179: Test Simple aggregator
-        agg = SimpleAggregator(lambda x, y: x + y)
-        m = ExternalMerger(agg, 20)
-        m.mergeValues(self.data)
-        self.assertTrue(m.spills >= 1)
-        self.assertEqual(sum(v for k, v in m.items()), sum(range(self.N)))
-
-    def test_merger_not_implemented_error(self):
-        # SPARK-39179: Test Merger for error scenarios
-        agg = SimpleAggregator(lambda x, y: x + y)
-
-        class DummyMerger(Merger):
-            def __init__(self, agg):
-                Merger.__init__(self, agg)
-
-        dummy_merger = DummyMerger(agg)
-        with self.assertRaises(NotImplementedError):
-            dummy_merger.mergeValues(self.data)
-        with self.assertRaises(NotImplementedError):
-            dummy_merger.mergeCombiners(self.data)
-        with self.assertRaises(NotImplementedError):
-            dummy_merger.items()
-
-    def test_huge_dataset(self):
-        m = ExternalMerger(self.agg, 5, partitions=3)
-        m.mergeCombiners(map(lambda k_v: (k_v[0], [str(k_v[1])]), self.data * 10))
-        self.assertTrue(m.spills >= 1)
-        self.assertEqual(sum(len(v) for k, v in m.items()), self.N * 10)
-        m._cleanup()
-
-    def test_group_by_key(self):
-        def gen_data(N, step):
-            for i in range(1, N + 1, step):
-                for j in range(i):
-                    yield (i, [j])
-
-        def gen_gs(N, step=1):
-            return shuffle.GroupByKey(gen_data(N, step))
-
-        self.assertEqual(1, len(list(gen_gs(1))))
-        self.assertEqual(2, len(list(gen_gs(2))))
-        self.assertEqual(100, len(list(gen_gs(100))))
-        self.assertEqual(list(range(1, 101)), [k for k, _ in gen_gs(100)])
-        self.assertTrue(all(list(range(k)) == list(vs) for k, vs in gen_gs(100)))
-
-        for k, vs in gen_gs(50002, 10000):
-            self.assertEqual(k, len(vs))
-            self.assertEqual(list(range(k)), list(vs))
-
-        ser = CPickleSerializer()
-        lst = ser.loads(ser.dumps(list(gen_gs(50002, 30000))))
-        for k, vs in lst:
-            self.assertEqual(k, len(vs))
-            self.assertEqual(list(range(k)), list(vs))
-
-    def test_stopiteration_is_raised(self):
-        def stopit(*args, **kwargs):
-            raise StopIteration()
-
-        def legit_create_combiner(x):
-            return [x]
-
-        def legit_merge_value(x, y):
-            return x.append(y) or x
-
-        def legit_merge_combiners(x, y):
-            return x.extend(y) or x
-
-        data = [(x % 2, x) for x in range(100)]
-
-        # wrong create combiner
-        m = ExternalMerger(Aggregator(stopit, legit_merge_value, legit_merge_combiners), 20)
-        with self.assertRaises((Py4JJavaError, RuntimeError)):
-            m.mergeValues(data)
-
-        # wrong merge value
-        m = ExternalMerger(Aggregator(legit_create_combiner, stopit, legit_merge_combiners), 20)
-        with self.assertRaises((Py4JJavaError, RuntimeError)):
-            m.mergeValues(data)
-
-        # wrong merge combiners
-        m = ExternalMerger(Aggregator(legit_create_combiner, legit_merge_value, stopit), 20)
-        with self.assertRaises((Py4JJavaError, RuntimeError)):
-            m.mergeCombiners(map(lambda x_y1: (x_y1[0], [x_y1[1]]), data))
-
-
-class ExternalGroupByTests(unittest.TestCase):
-    def setUp(self):
-        self.N = 1 << 20
-        values = [i for i in range(self.N)]
-        keys = [i for i in range(2)]
-        import itertools
-
-        self.data = [value for value in itertools.product(keys, values)]
-        self.agg = Aggregator(
-            lambda x: [x], lambda x, y: x.append(y) or x, lambda x, y: x.extend(y) or x
-        )
-
-    def test_medium_dataset(self):
-        # SPARK-39179: Test external group by for medium dataset
-        m = ExternalGroupBy(self.agg, 5, partitions=3)
-        m.mergeValues(self.data)
-        self.assertTrue(m.spills >= 1)
-        self.assertEqual(sum(sum(v) for k, v in m.items()), 2 * sum(range(self.N)))
-
-    def test_dataset_with_keys_are_unsorted(self):
-        # SPARK-39179: Test external group when numbers of keys are greater than SORT KEY Limit.
-        m = ExternalGroupBy(self.agg, 5, partitions=3)
-        original = m.SORT_KEY_LIMIT
-        try:
-            m.SORT_KEY_LIMIT = 1
-            m.mergeValues(self.data)
-            self.assertTrue(m.spills >= 1)
-            self.assertEqual(sum(sum(v) for k, v in m.items()), 2 * sum(range(self.N)))
-        finally:
-            m.SORT_KEY_LIMIT = original
-
-
-class SorterTests(unittest.TestCase):
-    def test_in_memory_sort(self):
-        lst = list(range(1024))
-        random.shuffle(lst)
-        sorter = ExternalSorter(1024)
-        self.assertEqual(sorted(lst), list(sorter.sorted(lst)))
-        self.assertEqual(sorted(lst, reverse=True), list(sorter.sorted(lst, reverse=True)))
-        self.assertEqual(sorted(lst, key=lambda x: -x), list(sorter.sorted(lst, key=lambda x: -x)))
-        self.assertEqual(
-            sorted(lst, key=lambda x: -x, reverse=True),
-            list(sorter.sorted(lst, key=lambda x: -x, reverse=True)),
-        )
-
-    def test_external_sort(self):
-        class CustomizedSorter(ExternalSorter):
-            def _next_limit(self):
-                return self.memory_limit
-
-        lst = list(range(1024))
-        random.shuffle(lst)
-        sorter = CustomizedSorter(1)
-        self.assertEqual(sorted(lst), list(sorter.sorted(lst)))
-        self.assertGreater(shuffle.DiskBytesSpilled, 0)
-        last = shuffle.DiskBytesSpilled
-        self.assertEqual(sorted(lst, reverse=True), list(sorter.sorted(lst, reverse=True)))
-        self.assertGreater(shuffle.DiskBytesSpilled, last)
-        last = shuffle.DiskBytesSpilled
-        self.assertEqual(sorted(lst, key=lambda x: -x), list(sorter.sorted(lst, key=lambda x: -x)))
-        self.assertGreater(shuffle.DiskBytesSpilled, last)
-        last = shuffle.DiskBytesSpilled
-        self.assertEqual(
-            sorted(lst, key=lambda x: -x, reverse=True),
-            list(sorter.sorted(lst, key=lambda x: -x, reverse=True)),
-        )
-        self.assertGreater(shuffle.DiskBytesSpilled, last)
-
-    def test_external_sort_in_rdd(self):
-        conf = SparkConf().set("spark.python.worker.memory", "1m")
-        sc = SparkContext(conf=conf)
-        lst = list(range(10240))
-        random.shuffle(lst)
-        rdd = sc.parallelize(lst, 4)
-        self.assertEqual(sorted(lst), rdd.sortBy(lambda x: x).collect())
-        sc.stop()
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_shuffle import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/tests/test_stage_sched.py

@@ -1,153 +0,0 @@
-#
-# 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
-import tempfile
-import unittest
-import time
-import shutil
-import json
-
-from pyspark import SparkConf, SparkContext
-from pyspark.resource.profile import ResourceProfileBuilder
-from pyspark.resource.requests import TaskResourceRequests
-from pyspark.taskcontext import TaskContext
-
-
-class StageSchedulingTest(unittest.TestCase):
-    def setUp(self):
-        self.temp_dir = tempfile.mkdtemp()
-
-    def tearDown(self):
-        shutil.rmtree(self.temp_dir, ignore_errors=True)
-        if getattr(self, "sc", None) is not None:
-            self.sc.stop()
-            self.sc = None
-
-    def _test_stage_scheduling(
-        self,
-        cpus_per_worker,
-        gpus_per_worker,
-        num_tasks,
-        resource_profile,
-        expected_max_concurrent_tasks,
-    ):
-        conf = SparkConf()
-        conf.setMaster(f"local-cluster[1,{cpus_per_worker},1024]").set(
-            "spark.task.maxFailures", "1"
-        )
-
-        if gpus_per_worker:
-            worker_res_config_file = os.path.join(self.temp_dir, "worker_res.json")
-            worker_res = [
-                {
-                    "id": {
-                        "componentName": "spark.worker",
-                        "resourceName": "gpu",
-                    },
-                    "addresses": [str(i) for i in range(gpus_per_worker)],
-                }
-            ]
-            with open(worker_res_config_file, "w") as fp:
-                json.dump(worker_res, fp)
-
-            conf.set("spark.worker.resource.gpu.amount", str(gpus_per_worker))
-            conf.set("spark.worker.resourcesFile", worker_res_config_file)
-            conf.set("spark.executor.resource.gpu.amount", str(gpus_per_worker))
-
-        self.sc = SparkContext(conf=conf)
-        pids_output_dir = os.path.join(self.temp_dir, "pids")
-        os.mkdir(pids_output_dir)
-
-        def mapper(_):
-            task_id = TaskContext.get().partitionId()
-            pid_file_path = os.path.join(pids_output_dir, str(task_id))
-            with open(pid_file_path, mode="w"):
-                pass
-            time.sleep(0.1)
-            num_concurrent_tasks = len(os.listdir(pids_output_dir))
-            time.sleep(1)
-            os.remove(pid_file_path)
-            return num_concurrent_tasks
-
-        results = (
-            self.sc.parallelize(range(num_tasks), num_tasks)
-            .withResources(resource_profile)
-            .map(mapper)
-            .collect()
-        )
-        self.assertEqual(max(results), expected_max_concurrent_tasks)
-
-    def test_stage_scheduling_3_cpu_per_task(self):
-        rp = ResourceProfileBuilder().require(TaskResourceRequests().cpus(3)).build
-        self._test_stage_scheduling(
-            cpus_per_worker=4,
-            gpus_per_worker=0,
-            num_tasks=2,
-            resource_profile=rp,
-            expected_max_concurrent_tasks=1,
-        )
-
-    def test_stage_scheduling_2_cpu_per_task(self):
-        rp = ResourceProfileBuilder().require(TaskResourceRequests().cpus(2)).build
-        self._test_stage_scheduling(
-            cpus_per_worker=4,
-            gpus_per_worker=0,
-            num_tasks=4,
-            resource_profile=rp,
-            expected_max_concurrent_tasks=2,
-        )
-
-    def test_stage_scheduling_2_cpus_2_gpus_per_task(self):
-        rp = (
-            ResourceProfileBuilder()
-            .require(TaskResourceRequests().cpus(2).resource("gpu", 2))
-            .build
-        )
-        self._test_stage_scheduling(
-            cpus_per_worker=4,
-            gpus_per_worker=4,
-            num_tasks=4,
-            resource_profile=rp,
-            expected_max_concurrent_tasks=2,
-        )
-
-    def test_stage_scheduling_2_cpus_3_gpus_per_task(self):
-        rp = (
-            ResourceProfileBuilder()
-            .require(TaskResourceRequests().cpus(2).resource("gpu", 3))
-            .build
-        )
-        self._test_stage_scheduling(
-            cpus_per_worker=4,
-            gpus_per_worker=4,
-            num_tasks=2,
-            resource_profile=rp,
-            expected_max_concurrent_tasks=1,
-        )
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_stage_sched import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)