concurrent.futures.ProcessPoolExecutor doesnt close all subprocesses when ctrl+C (KeyboardInterrupt) is used on Linux

[Little-Ryugu]

Bug report

The functions that use concurrent.futures.ProcessPoolExecutor have a bug when keyboard interrupt is used.

When a keyboard interrupt is used, the code (main process) errors out; however, subprocesses created from ProcessPoolExecutor do not error out and continue to run (called orphan processes, https://en.wikipedia.org/wiki/Orphan_process):

layup/src/layup/utilities/data_processing_utilities.py

Lines 239 to 247 in 6e8d197

	kwargs["primary_id_column_name"] = primary_id_column_name
	with ProcessPoolExecutor(max_workers=n_workers, mp_context=_MP_CONTEXT) as executor:
	# Create a future applying the function to each block of data for a given object id
	futures = [
	executor.submit(func, data[data[primary_id_column_name] == id], **kwargs)
	for id in np.unique(data[primary_id_column_name])
	]
	# Concatenate all processed blocks together as our final result
	return np.concatenate([future.result() for future in futures])

In edge cases where orbitfits conversion hangs and doesnt converge to an orbital solution, a user using ctrl+C (KeyboardInterrupt),might not realise that subprocesses are still running taking up computational resources.

A possible solution might be to use multiprocessing.Pool instead of ProcessPoolExecutor for parallel code so that we can terminate subprocesses on a keyboard interrupt, e.g https://noswap.com/blog/python-multiprocessing-keyboardinterrupt :

except KeyboardInterrupt:
        print "Caught KeyboardInterrupt, terminating workers"
        pool.terminate()
        pool.join()

layup/src/layup/utilities/data_processing_utilities.py

Lines 167 to 203 in 6e8d197

	def process_data(data, n_workers, func, **kwargs):
	"""
	Process a structured numpy array in parallel for a given function and keyword arguments
	Parameters
	----------
	data : numpy structured array
	The data to process.
	n_workers : int
	The number of workers to use for parallel processing.
	func : function
	The function to apply to each block of data within parallel.
	**kwargs : dictionary
	Extra arguments to pass to the function.
	Returns
	-------
	res : numpy structured array
	The processed data concatenated from each function result
	"""
	if n_workers < 1:
	raise ValueError(f"n_workers must be greater than 0, {n_workers} was provided.")
	if len(data) == 0:
	return data
	# Divide our data into blocks to be processed by each worker
	block_size = max(1, int(np.ceil(len(data) / n_workers)))
	# Create a list of tuples of the form (start, end) where start is the starting index of the block
	# and end is the last index of the block + 1.
	blocks = [(i, min(i + block_size, len(data))) for i in range(0, len(data), block_size)]
	with ProcessPoolExecutor(max_workers=n_workers, mp_context=_MP_CONTEXT) as executor:
	# Create a future applying the function to each block of data
	futures = [executor.submit(func, data[start:end], **kwargs) for start, end in blocks]
	# Concatenate all processed blocks together as our final result
	return np.concatenate([future.result() for future in futures])

layup/src/layup/utilities/data_processing_utilities.py

Lines 206 to 247 in 6e8d197

	def process_data_by_id(data, n_workers, func, primary_id_column_name, **kwargs):
	"""
	Process a structured numpy array in parallel for a given function and
	keyword arguments. Instead of distributing the data across all available workers
	it is expected that the data will contain a primary id column. The data will be
	split by the unique values in the primary id column and each block of data will
	be processed in parallel.
	Parameters
	----------
	data : numpy structured array
	The data to process. Expected to contain a primary id column.
	n_workers : int
	The number of workers to use for parallel processing.
	func : function
	The function to apply to each block of data within parallel.
	**kwargs : dictionary
	Extra arguments to pass to the function.
	Returns
	-------
	res : numpy structured array
	The processed data concatenated from each function result
	"""
	if n_workers < 1:
	raise ValueError(f"n_workers must be greater than 0, {n_workers} was provided.")
	#! Perhaps this should be None, or raise and exception that is caught by the
	#! caller. If we return `data`, the columns won't match the columns of the
	#! processed data.
	if len(data) == 0:
	return data
	kwargs["primary_id_column_name"] = primary_id_column_name
	with ProcessPoolExecutor(max_workers=n_workers, mp_context=_MP_CONTEXT) as executor:
	# Create a future applying the function to each block of data for a given object id
	futures = [
	executor.submit(func, data[data[primary_id_column_name] == id], **kwargs)
	for id in np.unique(data[primary_id_column_name])
	]
	# Concatenate all processed blocks together as our final result
	return np.concatenate([future.result() for future in futures])

Before submitting
Please check the following:

• I have described the situation in which the bug arose, including what code was executed, information about my environment, and any applicable data others will need to reproduce the problem.
• I have included available evidence of the unexpected behavior (including error messages, screenshots, and/or plots) as well as a description of what I expected instead.
• If I have a solution in mind, I have provided an explanation and/or pseudocode and/or task list.