Skip to content

Properly finalize ThreadPoolExecutors in file transfer agent #2879

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jochenott wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Properly finalize ThreadPoolExecutors in file transfer agent #2879

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 43 additions & 34 deletions src/snowflake/connector/file_transfer_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import sys
import threading
from concurrent.futures.thread import ThreadPoolExecutor
from contextlib import ExitStack
from dataclasses import dataclass
from functools import partial
from logging import getLogger
Expand Down Expand Up @@ -462,9 +463,16 @@ def transfer(self, metas: list[SnowflakeFileMeta]) -> None:
max_concurrency = min(
self._parallel, self._snowflake_server_dop_cap_for_file_transfer
)
network_tpe = ThreadPoolExecutor(max_concurrency)
preprocess_tpe = ThreadPoolExecutor(iobound_tpe_limit)
postprocess_tpe = ThreadPoolExecutor(iobound_tpe_limit)
thread_pools_exit_stack = ExitStack()
network_tpe = thread_pools_exit_stack.enter_context(
ThreadPoolExecutor(max_concurrency)
)
preprocess_tpe = thread_pools_exit_stack.enter_context(
ThreadPoolExecutor(iobound_tpe_limit)
)
postprocess_tpe = thread_pools_exit_stack.enter_context(
ThreadPoolExecutor(iobound_tpe_limit)
)
logger.debug(f"Chunk ThreadPoolExecutor size: {max_concurrency}")
cv_main_thread = threading.Condition() # to signal the main thread
cv_chunk_process = (
Expand Down Expand Up @@ -661,37 +669,38 @@ def function_and_callback_wrapper(
f"An exception was raised in {repr(callback)}", exc_info=True
)

for file_client in files:
callback = partial(preprocess_done_cb, done_client=file_client)
if is_upload:
preprocess_tpe.submit(
function_and_callback_wrapper,
# Work fn
file_client.prepare_upload,
# Callback fn
callback,
file_client.meta,
)
else:
preprocess_tpe.submit(
function_and_callback_wrapper,
# Work fn
file_client.prepare_download,
# Callback fn
callback,
file_client.meta,
)
transfer_metadata.num_files_started += 1 # TODO: do we need this?

with cv_main_thread:
while transfer_metadata.num_files_completed < num_total_files:
cv_main_thread.wait()
# If both exception_caught_in_work and exception_caught_in_callback
# are present, the former will take precedence.
if exception_caught_in_work is not None:
raise exception_caught_in_work
if exception_caught_in_callback is not None:
raise exception_caught_in_callback
with thread_pools_exit_stack:
for file_client in files:
callback = partial(preprocess_done_cb, done_client=file_client)
if is_upload:
preprocess_tpe.submit(
function_and_callback_wrapper,
# Work fn
file_client.prepare_upload,
# Callback fn
callback,
file_client.meta,
)
else:
preprocess_tpe.submit(
function_and_callback_wrapper,
# Work fn
file_client.prepare_download,
# Callback fn
callback,
file_client.meta,
)
transfer_metadata.num_files_started += 1 # TODO: do we need this?

with cv_main_thread:
while transfer_metadata.num_files_completed < num_total_files:
cv_main_thread.wait()
# If both exception_caught_in_work and exception_caught_in_callback
# are present, the former will take precedence.
if exception_caught_in_work is not None:
raise exception_caught_in_work
if exception_caught_in_callback is not None:
raise exception_caught_in_callback

self._results = metas

Expand Down
Loading