python/pyspark/tests/test_pin_thread.py - spark - Git at Google

 #
 # 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 time
 import threading
 import unittest

 from pyspark import SparkContext, SparkConf, InheritableThread


 class PinThreadTests(unittest.TestCase):
     # These tests are in a separate class because it uses
     # 'PYSPARK_PIN_THREAD' environment variable to test thread pin feature.

     @classmethod
     def setUpClass(cls):
         cls.old_pin_thread = os.environ.get("PYSPARK_PIN_THREAD")
         os.environ["PYSPARK_PIN_THREAD"] = "true"
         cls.sc = SparkContext("local[4]", cls.__name__, conf=SparkConf())

     @classmethod
     def tearDownClass(cls):
         cls.sc.stop()
         if cls.old_pin_thread is not None:
             os.environ["PYSPARK_PIN_THREAD"] = cls.old_pin_thread
         else:
             del os.environ["PYSPARK_PIN_THREAD"]

     def test_pinned_thread(self):
         threads = []
         exceptions = []
         property_name = "test_property_%s" % PinThreadTests.__name__
         jvm_thread_ids = []

         for i in range(10):

             def test_local_property():
                 jvm_thread_id = self.sc._jvm.java.lang.Thread.currentThread().getId()
                 jvm_thread_ids.append(jvm_thread_id)

                 # If a property is set in this thread, later it should get the same property
                 # within this thread.
                 self.sc.setLocalProperty(property_name, str(i))

                 # 5 threads, 1 second sleep. 5 threads without a sleep.
                 time.sleep(i % 2)

                 try:
                     assert self.sc.getLocalProperty(property_name) == str(i)

                     # Each command might create a thread in multi-threading mode in Py4J.
                     # This assert makes sure that the created thread is being reused.
                     assert jvm_thread_id == self.sc._jvm.java.lang.Thread.currentThread().getId()
                 except Exception as e:
                     exceptions.append(e)

             threads.append(threading.Thread(target=test_local_property))

         for t in threads:
             t.start()

         for t in threads:
             t.join()

         for e in exceptions:
             raise e

         # Created JVM threads should be 10 because Python thread are 10.
         assert len(set(jvm_thread_ids)) == 10

     def test_multiple_group_jobs(self):
         # SPARK-22340: Add a mode to pin Python thread into JVM's
         self.check_job_cancellation(
             lambda job_group: self.sc.setJobGroup(
                 job_group, "test rdd collect with setting job group"
             ),
             lambda job_group: self.sc.cancelJobGroup(job_group),
         )

     def test_multiple_group_tags(self):
         # SPARK-44194: Test pinned thread mode with job tags.
         self.check_job_cancellation(
             lambda job_tag: self.sc.addJobTag(job_tag),
             lambda job_tag: self.sc.cancelJobsWithTag(job_tag),
         )

     def check_job_cancellation(self, setter, canceller):
         job_id_a = "job_ids_to_cancel"
         job_id_b = "job_ids_to_run"

         threads = []
         thread_ids = range(4)
         thread_ids_to_cancel = [i for i in thread_ids if i % 2 == 0]
         thread_ids_to_run = [i for i in thread_ids if i % 2 != 0]

         # 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):
             """
             Executes a job with the group ``job_group``. Each job waits for 3 seconds
             and then exits.
             """
             try:
                 setter(job_id)
                 self.sc.parallelize([15]).map(lambda x: time.sleep(x)).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)
         self.assertFalse(is_job_cancelled[0])

         # 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)

         # Wait to make sure all jobs are executed.
         time.sleep(3)
         # And then, cancel one job group.
         canceller(job_id_a)

         # 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_local_property(self):
         self.sc.setLocalProperty("a", "hi")
         expected = []

         def get_inner_local_prop():
             expected.append(self.sc.getLocalProperty("b"))

         def get_outer_local_prop():
             expected.append(self.sc.getLocalProperty("a"))
             self.sc.setLocalProperty("b", "hello")
             t2 = InheritableThread(target=get_inner_local_prop)
             t2.start()
             t2.join()

         t1 = InheritableThread(target=get_outer_local_prop)
         t1.start()
         t1.join()

         self.assertEqual(self.sc.getLocalProperty("b"), None)
         self.assertEqual(expected, ["hi", "hello"])


 if __name__ == "__main__":
     import unittest
     from pyspark.tests.test_pin_thread 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)
	#
	# 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 time
	import threading
	import unittest

	from pyspark import SparkContext, SparkConf, InheritableThread


	class PinThreadTests(unittest.TestCase):
	# These tests are in a separate class because it uses
	# 'PYSPARK_PIN_THREAD' environment variable to test thread pin feature.

	@classmethod
	def setUpClass(cls):
	cls.old_pin_thread = os.environ.get("PYSPARK_PIN_THREAD")
	os.environ["PYSPARK_PIN_THREAD"] = "true"
	cls.sc = SparkContext("local[4]", cls.__name__, conf=SparkConf())

	@classmethod
	def tearDownClass(cls):
	cls.sc.stop()
	if cls.old_pin_thread is not None:
	os.environ["PYSPARK_PIN_THREAD"] = cls.old_pin_thread
	else:
	del os.environ["PYSPARK_PIN_THREAD"]

	def test_pinned_thread(self):
	threads = []
	exceptions = []
	property_name = "test_property_%s" % PinThreadTests.__name__
	jvm_thread_ids = []

	for i in range(10):

	def test_local_property():
	jvm_thread_id = self.sc._jvm.java.lang.Thread.currentThread().getId()
	jvm_thread_ids.append(jvm_thread_id)

	# If a property is set in this thread, later it should get the same property
	# within this thread.
	self.sc.setLocalProperty(property_name, str(i))

	# 5 threads, 1 second sleep. 5 threads without a sleep.
	time.sleep(i % 2)

	try:
	assert self.sc.getLocalProperty(property_name) == str(i)

	# Each command might create a thread in multi-threading mode in Py4J.
	# This assert makes sure that the created thread is being reused.
	assert jvm_thread_id == self.sc._jvm.java.lang.Thread.currentThread().getId()
	except Exception as e:
	exceptions.append(e)

	threads.append(threading.Thread(target=test_local_property))

	for t in threads:
	t.start()

	for t in threads:
	t.join()

	for e in exceptions:
	raise e

	# Created JVM threads should be 10 because Python thread are 10.
	assert len(set(jvm_thread_ids)) == 10

	def test_multiple_group_jobs(self):
	# SPARK-22340: Add a mode to pin Python thread into JVM's
	self.check_job_cancellation(
	lambda job_group: self.sc.setJobGroup(
	job_group, "test rdd collect with setting job group"
	),
	lambda job_group: self.sc.cancelJobGroup(job_group),
	)

	def test_multiple_group_tags(self):
	# SPARK-44194: Test pinned thread mode with job tags.
	self.check_job_cancellation(
	lambda job_tag: self.sc.addJobTag(job_tag),
	lambda job_tag: self.sc.cancelJobsWithTag(job_tag),
	)

	def check_job_cancellation(self, setter, canceller):
	job_id_a = "job_ids_to_cancel"
	job_id_b = "job_ids_to_run"

	threads = []
	thread_ids = range(4)
	thread_ids_to_cancel = [i for i in thread_ids if i % 2 == 0]
	thread_ids_to_run = [i for i in thread_ids if i % 2 != 0]

	# 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):
	"""
	Executes a job with the group ``job_group``. Each job waits for 3 seconds
	and then exits.
	"""
	try:
	setter(job_id)
	self.sc.parallelize([15]).map(lambda x: time.sleep(x)).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)
	self.assertFalse(is_job_cancelled[0])

	# 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)

	# Wait to make sure all jobs are executed.
	time.sleep(3)
	# And then, cancel one job group.
	canceller(job_id_a)

	# 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_local_property(self):
	self.sc.setLocalProperty("a", "hi")
	expected = []

	def get_inner_local_prop():
	expected.append(self.sc.getLocalProperty("b"))

	def get_outer_local_prop():
	expected.append(self.sc.getLocalProperty("a"))
	self.sc.setLocalProperty("b", "hello")
	t2 = InheritableThread(target=get_inner_local_prop)
	t2.start()
	t2.join()

	t1 = InheritableThread(target=get_outer_local_prop)
	t1.start()
	t1.join()

	self.assertEqual(self.sc.getLocalProperty("b"), None)
	self.assertEqual(expected, ["hi", "hello"])


	if __name__ == "__main__":
	import unittest
	from pyspark.tests.test_pin_thread 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)