codekingpro/portable-devtools
114k
1# Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License"). You4# may not use this file except in compliance with the License. A copy of5# the License is located at6#7# http://aws.amazon.com/apache2.0/8#9# or in the "license" file accompanying this file. This file is10# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF11# ANY KIND, either express or implied. See the License for the specific12# language governing permissions and limitations under the License.13"""Speeds up S3 throughput by using processes14 15Getting Started16===============17 18The :class:`ProcessPoolDownloader` can be used to download a single file by19calling :meth:`ProcessPoolDownloader.download_file`:20 21.. code:: python22 23 from s3transfer.processpool import ProcessPoolDownloader24 25 with ProcessPoolDownloader() as downloader:26 downloader.download_file('mybucket', 'mykey', 'myfile')27 28 29This snippet downloads the S3 object located in the bucket ``mybucket`` at the30key ``mykey`` to the local file ``myfile``. Any errors encountered during the31transfer are not propagated. To determine if a transfer succeeded or32failed, use the `Futures`_ interface.33 34 35The :class:`ProcessPoolDownloader` can be used to download multiple files as36well:37 38.. code:: python39 40 from s3transfer.processpool import ProcessPoolDownloader41 42 with ProcessPoolDownloader() as downloader:43 downloader.download_file('mybucket', 'mykey', 'myfile')44 downloader.download_file('mybucket', 'myotherkey', 'myotherfile')45 46 47When running this snippet, the downloading of ``mykey`` and ``myotherkey``48happen in parallel. The first ``download_file`` call does not block the49second ``download_file`` call. The snippet blocks when exiting50the context manager and blocks until both downloads are complete.51 52Alternatively, the ``ProcessPoolDownloader`` can be instantiated53and explicitly be shutdown using :meth:`ProcessPoolDownloader.shutdown`:54 55.. code:: python56 57 from s3transfer.processpool import ProcessPoolDownloader58 59 downloader = ProcessPoolDownloader()60 downloader.download_file('mybucket', 'mykey', 'myfile')61 downloader.download_file('mybucket', 'myotherkey', 'myotherfile')62 downloader.shutdown()63 64 65For this code snippet, the call to ``shutdown`` blocks until both66downloads are complete.67 68 69Additional Parameters70=====================71 72Additional parameters can be provided to the ``download_file`` method:73 74* ``extra_args``: A dictionary containing any additional client arguments75 to include in the76 `GetObject <https://botocore.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.get_object>`_77 API request. For example:78 79 .. code:: python80 81 from s3transfer.processpool import ProcessPoolDownloader82 83 with ProcessPoolDownloader() as downloader:84 downloader.download_file(85 'mybucket', 'mykey', 'myfile',86 extra_args={'VersionId': 'myversion'})87 88 89* ``expected_size``: By default, the downloader will make a HeadObject90 call to determine the size of the object. To opt-out of this additional91 API call, you can provide the size of the object in bytes:92 93 .. code:: python94 95 from s3transfer.processpool import ProcessPoolDownloader96 97 MB = 1024 * 102498 with ProcessPoolDownloader() as downloader:99 downloader.download_file(100 'mybucket', 'mykey', 'myfile', expected_size=2 * MB)101 102 103Futures104=======105 106When ``download_file`` is called, it immediately returns a107:class:`ProcessPoolTransferFuture`. The future can be used to poll the state108of a particular transfer. To get the result of the download,109call :meth:`ProcessPoolTransferFuture.result`. The method blocks110until the transfer completes, whether it succeeds or fails. For example:111 112.. code:: python113 114 from s3transfer.processpool import ProcessPoolDownloader115 116 with ProcessPoolDownloader() as downloader:117 future = downloader.download_file('mybucket', 'mykey', 'myfile')118 print(future.result())119 120 121If the download succeeds, the future returns ``None``:122 123.. code:: python124 125 None126 127 128If the download fails, the exception causing the failure is raised. For129example, if ``mykey`` did not exist, the following error would be raised130 131 132.. code:: python133 134 botocore.exceptions.ClientError: An error occurred (404) when calling the HeadObject operation: Not Found135 136 137.. note::138 139 :meth:`ProcessPoolTransferFuture.result` can only be called while the140 ``ProcessPoolDownloader`` is running (e.g. before calling ``shutdown`` or141 inside the context manager).142 143 144Process Pool Configuration145==========================146 147By default, the downloader has the following configuration options:148 149* ``multipart_threshold``: The threshold size for performing ranged downloads150 in bytes. By default, ranged downloads happen for S3 objects that are151 greater than or equal to 8 MB in size.152 153* ``multipart_chunksize``: The size of each ranged download in bytes. By154 default, the size of each ranged download is 8 MB.155 156* ``max_request_processes``: The maximum number of processes used to download157 S3 objects. By default, the maximum is 10 processes.158 159 160To change the default configuration, use the :class:`ProcessTransferConfig`:161 162.. code:: python163 164 from s3transfer.processpool import ProcessPoolDownloader165 from s3transfer.processpool import ProcessTransferConfig166 167 config = ProcessTransferConfig(168 multipart_threshold=64 * 1024 * 1024, # 64 MB169 max_request_processes=50170 )171 downloader = ProcessPoolDownloader(config=config)172 173 174Client Configuration175====================176 177The process pool downloader creates ``botocore`` clients on your behalf. In178order to affect how the client is created, pass the keyword arguments179that would have been used in the :meth:`botocore.Session.create_client` call:180 181.. code:: python182 183 184 from s3transfer.processpool import ProcessPoolDownloader185 from s3transfer.processpool import ProcessTransferConfig186 187 downloader = ProcessPoolDownloader(188 client_kwargs={'region_name': 'us-west-2'})189 190 191This snippet ensures that all clients created by the ``ProcessPoolDownloader``192are using ``us-west-2`` as their region.193 194"""195import collections196import contextlib197import logging198import multiprocessing199import signal200import threading201from copy import deepcopy202 203import botocore.session204from botocore.config import Config205 206from s3transfer.compat import MAXINT, BaseManager207from s3transfer.constants import ALLOWED_DOWNLOAD_ARGS, MB, PROCESS_USER_AGENT208from s3transfer.exceptions import CancelledError, RetriesExceededError209from s3transfer.futures import BaseTransferFuture, BaseTransferMeta210from s3transfer.utils import (211 S3_RETRYABLE_DOWNLOAD_ERRORS,212 CallArgs,213 OSUtils,214 calculate_num_parts,215 calculate_range_parameter,216)217 218logger = logging.getLogger(__name__)219 220SHUTDOWN_SIGNAL = 'SHUTDOWN'221 222# The DownloadFileRequest tuple is submitted from the ProcessPoolDownloader223# to the GetObjectSubmitter in order for the submitter to begin submitting224# GetObjectJobs to the GetObjectWorkers.225DownloadFileRequest = collections.namedtuple(226 'DownloadFileRequest',227 [228 'transfer_id', # The unique id for the transfer229 'bucket', # The bucket to download the object from230 'key', # The key to download the object from231 'filename', # The user-requested download location232 'extra_args', # Extra arguments to provide to client calls233 'expected_size', # The user-provided expected size of the download234 ],235)236 237# The GetObjectJob tuple is submitted from the GetObjectSubmitter238# to the GetObjectWorkers to download the file or parts of the file.239GetObjectJob = collections.namedtuple(240 'GetObjectJob',241 [242 'transfer_id', # The unique id for the transfer243 'bucket', # The bucket to download the object from244 'key', # The key to download the object from245 'temp_filename', # The temporary file to write the content to via246 # completed GetObject calls.247 'extra_args', # Extra arguments to provide to the GetObject call248 'offset', # The offset to write the content for the temp file.249 'filename', # The user-requested download location. The worker250 # of final GetObjectJob will move the file located at251 # temp_filename to the location of filename.252 ],253)254 255 256@contextlib.contextmanager257def ignore_ctrl_c():258 original_handler = _add_ignore_handler_for_interrupts()259 yield260 signal.signal(signal.SIGINT, original_handler)261 262 263def _add_ignore_handler_for_interrupts():264 # Windows is unable to pickle signal.signal directly so it needs to265 # be wrapped in a function defined at the module level266 return signal.signal(signal.SIGINT, signal.SIG_IGN)267 268 269class ProcessTransferConfig:270 def __init__(271 self,272 multipart_threshold=8 * MB,273 multipart_chunksize=8 * MB,274 max_request_processes=10,275 ):276 """Configuration for the ProcessPoolDownloader277 278 :param multipart_threshold: The threshold for which ranged downloads279 occur.280 281 :param multipart_chunksize: The chunk size of each ranged download.282 283 :param max_request_processes: The maximum number of processes that284 will be making S3 API transfer-related requests at a time.285 """286 self.multipart_threshold = multipart_threshold287 self.multipart_chunksize = multipart_chunksize288 self.max_request_processes = max_request_processes289 290 291class ProcessPoolDownloader:292 def __init__(self, client_kwargs=None, config=None):293 """Downloads S3 objects using process pools294 295 :type client_kwargs: dict296 :param client_kwargs: The keyword arguments to provide when297 instantiating S3 clients. The arguments must match the keyword298 arguments provided to the299 `botocore.session.Session.create_client()` method.300 301 :type config: ProcessTransferConfig302 :param config: Configuration for the downloader303 """304 if client_kwargs is None:305 client_kwargs = {}306 self._client_factory = ClientFactory(client_kwargs)307 308 self._transfer_config = config309 if config is None:310 self._transfer_config = ProcessTransferConfig()311 312 self._download_request_queue = multiprocessing.Queue(1000)313 self._worker_queue = multiprocessing.Queue(1000)314 self._osutil = OSUtils()315 316 self._started = False317 self._start_lock = threading.Lock()318 319 # These below are initialized in the start() method320 self._manager = None321 self._transfer_monitor = None322 self._submitter = None323 self._workers = []324 325 def download_file(326 self, bucket, key, filename, extra_args=None, expected_size=None327 ):328 """Downloads the object's contents to a file329 330 :type bucket: str331 :param bucket: The name of the bucket to download from332 333 :type key: str334 :param key: The name of the key to download from335 336 :type filename: str337 :param filename: The name of a file to download to.338 339 :type extra_args: dict340 :param extra_args: Extra arguments that may be passed to the341 client operation342 343 :type expected_size: int344 :param expected_size: The expected size in bytes of the download. If345 provided, the downloader will not call HeadObject to determine the346 object's size and use the provided value instead. The size is347 needed to determine whether to do a multipart download.348 349 :rtype: s3transfer.futures.TransferFuture350 :returns: Transfer future representing the download351 """352 self._start_if_needed()353 if extra_args is None:354 extra_args = {}355 self._validate_all_known_args(extra_args)356 transfer_id = self._transfer_monitor.notify_new_transfer()357 download_file_request = DownloadFileRequest(358 transfer_id=transfer_id,359 bucket=bucket,360 key=key,361 filename=filename,362 extra_args=extra_args,363 expected_size=expected_size,364 )365 logger.debug(366 'Submitting download file request: %s.', download_file_request367 )368 self._download_request_queue.put(download_file_request)369 call_args = CallArgs(370 bucket=bucket,371 key=key,372 filename=filename,373 extra_args=extra_args,374 expected_size=expected_size,375 )376 future = self._get_transfer_future(transfer_id, call_args)377 return future378 379 def shutdown(self):380 """Shutdown the downloader381 382 It will wait till all downloads are complete before returning.383 """384 self._shutdown_if_needed()385 386 def __enter__(self):387 return self388 389 def __exit__(self, exc_type, exc_value, *args):390 if isinstance(exc_value, KeyboardInterrupt):391 if self._transfer_monitor is not None:392 self._transfer_monitor.notify_cancel_all_in_progress()393 self.shutdown()394 395 def _start_if_needed(self):396 with self._start_lock:397 if not self._started:398 self._start()399 400 def _start(self):401 self._start_transfer_monitor_manager()402 self._start_submitter()403 self._start_get_object_workers()404 self._started = True405 406 def _validate_all_known_args(self, provided):407 for kwarg in provided:408 if kwarg not in ALLOWED_DOWNLOAD_ARGS:409 download_args = ', '.join(ALLOWED_DOWNLOAD_ARGS)410 raise ValueError(411 f"Invalid extra_args key '{kwarg}', "412 f"must be one of: {download_args}"413 )414 415 def _get_transfer_future(self, transfer_id, call_args):416 meta = ProcessPoolTransferMeta(417 call_args=call_args, transfer_id=transfer_id418 )419 future = ProcessPoolTransferFuture(420 monitor=self._transfer_monitor, meta=meta421 )422 return future423 424 def _start_transfer_monitor_manager(self):425 logger.debug('Starting the TransferMonitorManager.')426 self._manager = TransferMonitorManager()427 # We do not want Ctrl-C's to cause the manager to shutdown immediately428 # as worker processes will still need to communicate with it when they429 # are shutting down. So instead we ignore Ctrl-C and let the manager430 # be explicitly shutdown when shutting down the downloader.431 self._manager.start(_add_ignore_handler_for_interrupts)432 self._transfer_monitor = self._manager.TransferMonitor()433 434 def _start_submitter(self):435 logger.debug('Starting the GetObjectSubmitter.')436 self._submitter = GetObjectSubmitter(437 transfer_config=self._transfer_config,438 client_factory=self._client_factory,439 transfer_monitor=self._transfer_monitor,440 osutil=self._osutil,441 download_request_queue=self._download_request_queue,442 worker_queue=self._worker_queue,443 )444 self._submitter.start()445 446 def _start_get_object_workers(self):447 logger.debug(448 'Starting %s GetObjectWorkers.',449 self._transfer_config.max_request_processes,450 )451 for _ in range(self._transfer_config.max_request_processes):452 worker = GetObjectWorker(453 queue=self._worker_queue,454 client_factory=self._client_factory,455 transfer_monitor=self._transfer_monitor,456 osutil=self._osutil,457 )458 worker.start()459 self._workers.append(worker)460 461 def _shutdown_if_needed(self):462 with self._start_lock:463 if self._started:464 self._shutdown()465 466 def _shutdown(self):467 self._shutdown_submitter()468 self._shutdown_get_object_workers()469 self._shutdown_transfer_monitor_manager()470 self._started = False471 472 def _shutdown_transfer_monitor_manager(self):473 logger.debug('Shutting down the TransferMonitorManager.')474 self._manager.shutdown()475 476 def _shutdown_submitter(self):477 logger.debug('Shutting down the GetObjectSubmitter.')478 self._download_request_queue.put(SHUTDOWN_SIGNAL)479 self._submitter.join()480 481 def _shutdown_get_object_workers(self):482 logger.debug('Shutting down the GetObjectWorkers.')483 for _ in self._workers:484 self._worker_queue.put(SHUTDOWN_SIGNAL)485 for worker in self._workers:486 worker.join()487 488 489class ProcessPoolTransferFuture(BaseTransferFuture):490 def __init__(self, monitor, meta):491 """The future associated to a submitted process pool transfer request492 493 :type monitor: TransferMonitor494 :param monitor: The monitor associated to the process pool downloader495 496 :type meta: ProcessPoolTransferMeta497 :param meta: The metadata associated to the request. This object498 is visible to the requester.499 """500 self._monitor = monitor501 self._meta = meta502 503 @property504 def meta(self):505 return self._meta506 507 def done(self):508 return self._monitor.is_done(self._meta.transfer_id)509 510 def result(self):511 try:512 return self._monitor.poll_for_result(self._meta.transfer_id)513 except KeyboardInterrupt:514 # For the multiprocessing Manager, a thread is given a single515 # connection to reuse in communicating between the thread in the516 # main process and the Manager's process. If a Ctrl-C happens when517 # polling for the result, it will make the main thread stop trying518 # to receive from the connection, but the Manager process will not519 # know that the main process has stopped trying to receive and520 # will not close the connection. As a result if another message is521 # sent to the Manager process, the listener in the Manager522 # processes will not process the new message as it is still trying523 # trying to process the previous message (that was Ctrl-C'd) and524 # thus cause the thread in the main process to hang on its send.525 # The only way around this is to create a new connection and send526 # messages from that new connection instead.527 self._monitor._connect()528 self.cancel()529 raise530 531 def cancel(self):532 self._monitor.notify_exception(533 self._meta.transfer_id, CancelledError()534 )535 536 537class ProcessPoolTransferMeta(BaseTransferMeta):538 """Holds metadata about the ProcessPoolTransferFuture"""539 540 def __init__(self, transfer_id, call_args):541 self._transfer_id = transfer_id542 self._call_args = call_args543 self._user_context = {}544 545 @property546 def call_args(self):547 return self._call_args548 549 @property550 def transfer_id(self):551 return self._transfer_id552 553 @property554 def user_context(self):555 return self._user_context556 557 558class ClientFactory:559 def __init__(self, client_kwargs=None):560 """Creates S3 clients for processes561 562 Botocore sessions and clients are not pickleable so they cannot be563 inherited across Process boundaries. Instead, they must be instantiated564 once a process is running.565 """566 self._client_kwargs = client_kwargs567 if self._client_kwargs is None:568 self._client_kwargs = {}569 570 client_config = deepcopy(self._client_kwargs.get('config', Config()))571 if not client_config.user_agent_extra:572 client_config.user_agent_extra = PROCESS_USER_AGENT573 else:574 client_config.user_agent_extra += " " + PROCESS_USER_AGENT575 self._client_kwargs['config'] = client_config576 577 def create_client(self):578 """Create a botocore S3 client"""579 return botocore.session.Session().create_client(580 's3', **self._client_kwargs581 )582 583 584class TransferMonitor:585 def __init__(self):586 """Monitors transfers for cross-process communication587 588 Notifications can be sent to the monitor and information can be589 retrieved from the monitor for a particular transfer. This abstraction590 is ran in a ``multiprocessing.managers.BaseManager`` in order to be591 shared across processes.592 """593 # TODO: Add logic that removes the TransferState if the transfer is594 # marked as done and the reference to the future is no longer being595 # held onto. Without this logic, this dictionary will continue to596 # grow in size with no limit.597 self._transfer_states = {}598 self._id_count = 0599 self._init_lock = threading.Lock()600 601 def notify_new_transfer(self):602 with self._init_lock:603 transfer_id = self._id_count604 self._transfer_states[transfer_id] = TransferState()605 self._id_count += 1606 return transfer_id607 608 def is_done(self, transfer_id):609 """Determine a particular transfer is complete610 611 :param transfer_id: Unique identifier for the transfer612 :return: True, if done. False, otherwise.613 """614 return self._transfer_states[transfer_id].done615 616 def notify_done(self, transfer_id):617 """Notify a particular transfer is complete618 619 :param transfer_id: Unique identifier for the transfer620 """621 self._transfer_states[transfer_id].set_done()622 623 def poll_for_result(self, transfer_id):624 """Poll for the result of a transfer625 626 :param transfer_id: Unique identifier for the transfer627 :return: If the transfer succeeded, it will return the result. If the628 transfer failed, it will raise the exception associated to the629 failure.630 """631 self._transfer_states[transfer_id].wait_till_done()632 exception = self._transfer_states[transfer_id].exception633 if exception:634 raise exception635 return None636 637 def notify_exception(self, transfer_id, exception):638 """Notify an exception was encountered for a transfer639 640 :param transfer_id: Unique identifier for the transfer641 :param exception: The exception encountered for that transfer642 """643 # TODO: Not all exceptions are pickleable so if we are running644 # this in a multiprocessing.BaseManager we will want to645 # make sure to update this signature to ensure pickleability of the646 # arguments or have the ProxyObject do the serialization.647 self._transfer_states[transfer_id].exception = exception648 649 def notify_cancel_all_in_progress(self):650 for transfer_state in self._transfer_states.values():651 if not transfer_state.done:652 transfer_state.exception = CancelledError()653 654 def get_exception(self, transfer_id):655 """Retrieve the exception encountered for the transfer656 657 :param transfer_id: Unique identifier for the transfer658 :return: The exception encountered for that transfer. Otherwise659 if there were no exceptions, returns None.660 """661 return self._transfer_states[transfer_id].exception662 663 def notify_expected_jobs_to_complete(self, transfer_id, num_jobs):664 """Notify the amount of jobs expected for a transfer665 666 :param transfer_id: Unique identifier for the transfer667 :param num_jobs: The number of jobs to complete the transfer668 """669 self._transfer_states[transfer_id].jobs_to_complete = num_jobs670 671 def notify_job_complete(self, transfer_id):672 """Notify that a single job is completed for a transfer673 674 :param transfer_id: Unique identifier for the transfer675 :return: The number of jobs remaining to complete the transfer676 """677 return self._transfer_states[transfer_id].decrement_jobs_to_complete()678 679 680class TransferState:681 """Represents the current state of an individual transfer"""682 683 # NOTE: Ideally the TransferState object would be used directly by the684 # various different abstractions in the ProcessPoolDownloader and remove685 # the need for the TransferMonitor. However, it would then impose the686 # constraint that two hops are required to make or get any changes in the687 # state of a transfer across processes: one hop to get a proxy object for688 # the TransferState and then a second hop to communicate calling the689 # specific TransferState method.690 def __init__(self):691 self._exception = None692 self._done_event = threading.Event()693 self._job_lock = threading.Lock()694 self._jobs_to_complete = 0695 696 @property697 def done(self):698 return self._done_event.is_set()699 700 def set_done(self):701 self._done_event.set()702 703 def wait_till_done(self):704 self._done_event.wait(MAXINT)705 706 @property707 def exception(self):708 return self._exception709 710 @exception.setter711 def exception(self, val):712 self._exception = val713 714 @property715 def jobs_to_complete(self):716 return self._jobs_to_complete717 718 @jobs_to_complete.setter719 def jobs_to_complete(self, val):720 self._jobs_to_complete = val721 722 def decrement_jobs_to_complete(self):723 with self._job_lock:724 self._jobs_to_complete -= 1725 return self._jobs_to_complete726 727 728class TransferMonitorManager(BaseManager):729 pass730 731 732TransferMonitorManager.register('TransferMonitor', TransferMonitor)733 734 735class BaseS3TransferProcess(multiprocessing.Process):736 def __init__(self, client_factory):737 super().__init__()738 self._client_factory = client_factory739 self._client = None740 741 def run(self):742 # Clients are not pickleable so their instantiation cannot happen743 # in the __init__ for processes that are created under the744 # spawn method.745 self._client = self._client_factory.create_client()746 with ignore_ctrl_c():747 # By default these processes are ran as child processes to the748 # main process. Any Ctrl-c encountered in the main process is749 # propagated to the child process and interrupt it at any time.750 # To avoid any potentially bad states caused from an interrupt751 # (i.e. a transfer failing to notify its done or making the752 # communication protocol become out of sync with the753 # TransferMonitor), we ignore all Ctrl-C's and allow the main754 # process to notify these child processes when to stop processing755 # jobs.756 self._do_run()757 758 def _do_run(self):759 raise NotImplementedError('_do_run()')760 761 762class GetObjectSubmitter(BaseS3TransferProcess):763 def __init__(764 self,765 transfer_config,766 client_factory,767 transfer_monitor,768 osutil,769 download_request_queue,770 worker_queue,771 ):772 """Submit GetObjectJobs to fulfill a download file request773 774 :param transfer_config: Configuration for transfers.775 :param client_factory: ClientFactory for creating S3 clients.776 :param transfer_monitor: Monitor for notifying and retrieving state777 of transfer.778 :param osutil: OSUtils object to use for os-related behavior when779 performing the transfer.780 :param download_request_queue: Queue to retrieve download file781 requests.782 :param worker_queue: Queue to submit GetObjectJobs for workers783 to perform.784 """785 super().__init__(client_factory)786 self._transfer_config = transfer_config787 self._transfer_monitor = transfer_monitor788 self._osutil = osutil789 self._download_request_queue = download_request_queue790 self._worker_queue = worker_queue791 792 def _do_run(self):793 while True:794 download_file_request = self._download_request_queue.get()795 if download_file_request == SHUTDOWN_SIGNAL:796 logger.debug('Submitter shutdown signal received.')797 return798 try:799 self._submit_get_object_jobs(download_file_request)800 except Exception as e:801 logger.debug(802 'Exception caught when submitting jobs for '803 'download file request %s: %s',804 download_file_request,805 e,806 exc_info=True,807 )808 self._transfer_monitor.notify_exception(809 download_file_request.transfer_id, e810 )811 self._transfer_monitor.notify_done(812 download_file_request.transfer_id813 )814 815 def _submit_get_object_jobs(self, download_file_request):816 size = self._get_size(download_file_request)817 temp_filename = self._allocate_temp_file(download_file_request, size)818 if size < self._transfer_config.multipart_threshold:819 self._submit_single_get_object_job(820 download_file_request, temp_filename821 )822 else:823 self._submit_ranged_get_object_jobs(824 download_file_request, temp_filename, size825 )826 827 def _get_size(self, download_file_request):828 expected_size = download_file_request.expected_size829 if expected_size is None:830 expected_size = self._client.head_object(831 Bucket=download_file_request.bucket,832 Key=download_file_request.key,833 **download_file_request.extra_args,834 )['ContentLength']835 return expected_size836 837 def _allocate_temp_file(self, download_file_request, size):838 temp_filename = self._osutil.get_temp_filename(839 download_file_request.filename840 )841 self._osutil.allocate(temp_filename, size)842 return temp_filename843 844 def _submit_single_get_object_job(845 self, download_file_request, temp_filename846 ):847 self._notify_jobs_to_complete(download_file_request.transfer_id, 1)848 self._submit_get_object_job(849 transfer_id=download_file_request.transfer_id,850 bucket=download_file_request.bucket,851 key=download_file_request.key,852 temp_filename=temp_filename,853 offset=0,854 extra_args=download_file_request.extra_args,855 filename=download_file_request.filename,856 )857 858 def _submit_ranged_get_object_jobs(859 self, download_file_request, temp_filename, size860 ):861 part_size = self._transfer_config.multipart_chunksize862 num_parts = calculate_num_parts(size, part_size)863 self._notify_jobs_to_complete(864 download_file_request.transfer_id, num_parts865 )866 for i in range(num_parts):867 offset = i * part_size868 range_parameter = calculate_range_parameter(869 part_size, i, num_parts870 )871 get_object_kwargs = {'Range': range_parameter}872 get_object_kwargs.update(download_file_request.extra_args)873 self._submit_get_object_job(874 transfer_id=download_file_request.transfer_id,875 bucket=download_file_request.bucket,876 key=download_file_request.key,877 temp_filename=temp_filename,878 offset=offset,879 extra_args=get_object_kwargs,880 filename=download_file_request.filename,881 )882 883 def _submit_get_object_job(self, **get_object_job_kwargs):884 self._worker_queue.put(GetObjectJob(**get_object_job_kwargs))885 886 def _notify_jobs_to_complete(self, transfer_id, jobs_to_complete):887 logger.debug(888 'Notifying %s job(s) to complete for transfer_id %s.',889 jobs_to_complete,890 transfer_id,891 )892 self._transfer_monitor.notify_expected_jobs_to_complete(893 transfer_id, jobs_to_complete894 )895 896 897class GetObjectWorker(BaseS3TransferProcess):898 # TODO: It may make sense to expose these class variables as configuration899 # options if users want to tweak them.900 _MAX_ATTEMPTS = 5901 _IO_CHUNKSIZE = 2 * MB902 903 def __init__(self, queue, client_factory, transfer_monitor, osutil):904 """Fulfills GetObjectJobs905 906 Downloads the S3 object, writes it to the specified file, and907 renames the file to its final location if it completes the final908 job for a particular transfer.909 910 :param queue: Queue for retrieving GetObjectJob's911 :param client_factory: ClientFactory for creating S3 clients912 :param transfer_monitor: Monitor for notifying913 :param osutil: OSUtils object to use for os-related behavior when914 performing the transfer.915 """916 super().__init__(client_factory)917 self._queue = queue918 self._client_factory = client_factory919 self._transfer_monitor = transfer_monitor920 self._osutil = osutil921 922 def _do_run(self):923 while True:924 job = self._queue.get()925 if job == SHUTDOWN_SIGNAL:926 logger.debug('Worker shutdown signal received.')927 return928 if not self._transfer_monitor.get_exception(job.transfer_id):929 self._run_get_object_job(job)930 else:931 logger.debug(932 'Skipping get object job %s because there was a previous '933 'exception.',934 job,935 )936 remaining = self._transfer_monitor.notify_job_complete(937 job.transfer_id938 )939 logger.debug(940 '%s jobs remaining for transfer_id %s.',941 remaining,942 job.transfer_id,943 )944 if not remaining:945 self._finalize_download(946 job.transfer_id, job.temp_filename, job.filename947 )948 949 def _run_get_object_job(self, job):950 try:951 self._do_get_object(952 bucket=job.bucket,953 key=job.key,954 temp_filename=job.temp_filename,955 extra_args=job.extra_args,956 offset=job.offset,957 )958 except Exception as e:959 logger.debug(960 'Exception caught when downloading object for '961 'get object job %s: %s',962 job,963 e,964 exc_info=True,965 )966 self._transfer_monitor.notify_exception(job.transfer_id, e)967 968 def _do_get_object(self, bucket, key, extra_args, temp_filename, offset):969 last_exception = None970 for i in range(self._MAX_ATTEMPTS):971 try:972 response = self._client.get_object(973 Bucket=bucket, Key=key, **extra_args974 )975 self._write_to_file(temp_filename, offset, response['Body'])976 return977 except S3_RETRYABLE_DOWNLOAD_ERRORS as e:978 logger.debug(979 'Retrying exception caught (%s), '980 'retrying request, (attempt %s / %s)',981 e,982 i + 1,983 self._MAX_ATTEMPTS,984 exc_info=True,985 )986 last_exception = e987 raise RetriesExceededError(last_exception)988 989 def _write_to_file(self, filename, offset, body):990 with open(filename, 'rb+') as f:991 f.seek(offset)992 chunks = iter(lambda: body.read(self._IO_CHUNKSIZE), b'')993 for chunk in chunks:994 f.write(chunk)995 996 def _finalize_download(self, transfer_id, temp_filename, filename):997 if self._transfer_monitor.get_exception(transfer_id):998 self._osutil.remove_file(temp_filename)999 else:1000 self._do_file_rename(transfer_id, temp_filename, filename)1001 self._transfer_monitor.notify_done(transfer_id)1002 1003 def _do_file_rename(self, transfer_id, temp_filename, filename):1004 try:1005 self._osutil.rename_file(temp_filename, filename)1006 except Exception as e:1007 self._transfer_monitor.notify_exception(transfer_id, e)1008 self._osutil.remove_file(temp_filename)1009