feat: add transfer_manager.upload_chunks_concurrently using the XML MPU API (#1115)

* intermediate commit * temporary commit * xml mpu support, unit tests and docstrings * integration tests * add support for metadata * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * encryption support * unit tests for mpu * docs update * fix unit test issue * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: andrewsg

google/cloud/storage/blob.py

@@ -1697,7 +1697,7 @@ def _get_writable_metadata(self):
 
  return object_metadata
 
- def _get_upload_arguments(self, client, content_type):
+ def _get_upload_arguments(self, client, content_type, filename=None):
  """Get required arguments for performing an upload.
 
  The content type returned will be determined in order of precedence:
@@ -1716,7 +1716,7 @@ def _get_upload_arguments(self, client, content_type):
  * An object metadata dictionary
  * The ``content_type`` as a string (according to precedence)
  """
- content_type = self._get_content_type(content_type)
+ content_type = self._get_content_type(content_type, filename=filename)
  headers = {
  **_get_default_headers(client._connection.user_agent, content_type),
  **_get_encryption_headers(self._encryption_key),

google/cloud/storage/transfer_manager.py

@@ -26,6 +26,12 @@
 from google.api_core import exceptions
 from google.cloud.storage import Client
 from google.cloud.storage import Blob
+from google.cloud.storage.blob import _get_host_name
+from google.cloud.storage.constants import _DEFAULT_TIMEOUT
+
+from google.resumable_media.requests.upload import XMLMPUContainer
+from google.resumable_media.requests.upload import XMLMPUPart
+
 
 warnings.warn(
  "The module `transfer_manager` is a preview feature. Functionality and API "
@@ -35,7 +41,14 @@
 
 TM_DEFAULT_CHUNK_SIZE = 32 * 1024 * 1024
 DEFAULT_MAX_WORKERS = 8
-
+METADATA_HEADER_TRANSLATION = {
+ "cacheControl": "Cache-Control",
+ "contentDisposition": "Content-Disposition",
+ "contentEncoding": "Content-Encoding",
+ "contentLanguage": "Content-Language",
+ "customTime": "x-goog-custom-time",
+ "storageClass": "x-goog-storage-class",
+}
 
 # Constants to be passed in as `worker_type`.
 PROCESS = "process"
@@ -198,7 +211,7 @@ def upload_many(
  futures.append(
  executor.submit(
  _call_method_on_maybe_pickled_blob,
- _pickle_blob(blob) if needs_pickling else blob,
+ _pickle_client(blob) if needs_pickling else blob,
  "upload_from_filename"
  if isinstance(path_or_file, str)
  else "upload_from_file",
@@ -343,7 +356,7 @@ def download_many(
  futures.append(
  executor.submit(
  _call_method_on_maybe_pickled_blob,
- _pickle_blob(blob) if needs_pickling else blob,
+ _pickle_client(blob) if needs_pickling else blob,
  "download_to_filename"
  if isinstance(path_or_file, str)
  else "download_to_file",
@@ -733,7 +746,6 @@ def download_chunks_concurrently(
  Checksumming (md5 or crc32c) is not supported for chunked operations. Any
  `checksum` parameter passed in to download_kwargs will be ignored.
 
- :type bucket: 'google.cloud.storage.bucket.Bucket'
  :param bucket:
  The bucket which contains the blobs to be downloaded
 
@@ -745,6 +757,12 @@ def download_chunks_concurrently(
  :param filename:
  The destination filename or path.
 
+ :type chunk_size: int
+ :param chunk_size:
+ The size in bytes of each chunk to send. The optimal chunk size for
+ maximum throughput may vary depending on the exact network environment
+ and size of the blob.
+
  :type download_kwargs: dict
  :param download_kwargs:
  A dictionary of keyword arguments to pass to the download method. Refer
@@ -809,7 +827,7 @@ def download_chunks_concurrently(
 
  pool_class, needs_pickling = _get_pool_class_and_requirements(worker_type)
  # Pickle the blob ahead of time (just once, not once per chunk) if needed.
- maybe_pickled_blob = _pickle_blob(blob) if needs_pickling else blob
+ maybe_pickled_blob = _pickle_client(blob) if needs_pickling else blob
 
  futures = []
 
@@ -844,9 +862,249 @@ def download_chunks_concurrently(
  return None
 
 
+def upload_chunks_concurrently(
+ filename,
+ blob,
+ content_type=None,
+ chunk_size=TM_DEFAULT_CHUNK_SIZE,
+ deadline=None,
+ worker_type=PROCESS,
+ max_workers=DEFAULT_MAX_WORKERS,
+ *,
+ checksum="md5",
+ timeout=_DEFAULT_TIMEOUT,
+):
+ """Upload a single file in chunks, concurrently.
+
+ This function uses the XML MPU API to initialize an upload and upload a
+ file in chunks, concurrently with a worker pool.
+
+ The XML MPU API is significantly different from other uploads; please review
+ the documentation at https://cloud.google.com/storage/docs/multipart-uploads
+ before using this feature.
+
+ The library will attempt to cancel uploads that fail due to an exception.
+ If the upload fails in a way that precludes cancellation, such as a
+ hardware failure, process termination, or power outage, then the incomplete
+ upload may persist indefinitely. To mitigate this, set the
+ `AbortIncompleteMultipartUpload` with a nonzero `Age` in bucket lifecycle
+ rules, or refer to the XML API documentation linked above to learn more
+ about how to list and delete individual downloads.
+
+ Using this feature with multiple threads is unlikely to improve upload
+ performance under normal circumstances due to Python interpreter threading
+ behavior. The default is therefore to use processes instead of threads.
+
+ ACL information cannot be sent with this function and should be set
+ separately with :class:`ObjectACL` methods.
+
+ :type filename: str
+ :param filename:
+ The path to the file to upload. File-like objects are not supported.
+
+ :type blob: `google.cloud.storage.Blob`
+ :param blob:
+ The blob to which to upload.
+
+ :type content_type: str
+ :param content_type: (Optional) Type of content being uploaded.
+
+ :type chunk_size: int
+ :param chunk_size:
+ The size in bytes of each chunk to send. The optimal chunk size for
+ maximum throughput may vary depending on the exact network environment
+ and size of the blob. The remote API has restrictions on the minimum
+ and maximum size allowable, see: https://cloud.google.com/storage/quotas#requests
+
+ :type deadline: int
+ :param deadline:
+ The number of seconds to wait for all threads to resolve. If the
+ deadline is reached, all threads will be terminated regardless of their
+ progress and concurrent.futures.TimeoutError will be raised. This can be
+ left as the default of None (no deadline) for most use cases.
+
+ :type worker_type: str
+ :param worker_type:
+ The worker type to use; one of google.cloud.storage.transfer_manager.PROCESS
+ or google.cloud.storage.transfer_manager.THREAD.
+
+ Although the exact performance impact depends on the use case, in most
+ situations the PROCESS worker type will use more system resources (both
+ memory and CPU) and result in faster operations than THREAD workers.
+
+ Because the subprocesses of the PROCESS worker type can't access memory
+ from the main process, Client objects have to be serialized and then
+ recreated in each subprocess. The serialization of the Client object
+ for use in subprocesses is an approximation and may not capture every
+ detail of the Client object, especially if the Client was modified after
+ its initial creation or if `Client._http` was modified in any way.
+
+ THREAD worker types are observed to be relatively efficient for
+ operations with many small files, but not for operations with large
+ files. PROCESS workers are recommended for large file operations.
+
+ :type max_workers: int
+ :param max_workers:
+ The maximum number of workers to create to handle the workload.
+
+ With PROCESS workers, a larger number of workers will consume more
+ system resources (memory and CPU) at once.
+
+ How many workers is optimal depends heavily on the specific use case,
+ and the default is a conservative number that should work okay in most
+ cases without consuming excessive resources.
+
+ :type checksum: str
+ :param checksum:
+ (Optional) The checksum scheme to use: either 'md5', 'crc32c' or None.
+ Each individual part is checksummed. At present, the selected checksum
+ rule is only applied to parts and a separate checksum of the entire
+ resulting blob is not computed. Please compute and compare the checksum
+ of the file to the resulting blob separately if needed, using the
+ 'crc32c' algorithm as per the XML MPU documentation.
+
+ :type timeout: float or tuple
+ :param timeout:
+ (Optional) The amount of time, in seconds, to wait
+ for the server response. See: :ref:`configuring_timeouts`
+
+ :raises: :exc:`concurrent.futures.TimeoutError` if deadline is exceeded.
+ """
+
+ bucket = blob.bucket
+ client = blob.client
+ transport = blob._get_transport(client)
+
+ hostname = _get_host_name(client._connection)
+ url = "{hostname}/{bucket}/{blob}".format(
+ hostname=hostname, bucket=bucket.name, blob=blob.name
+ )
+
+ base_headers, object_metadata, content_type = blob._get_upload_arguments(
+ client, content_type, filename=filename
+ )
+ headers = {**base_headers, **_headers_from_metadata(object_metadata)}
+
+ if blob.user_project is not None:
+ headers["x-goog-user-project"] = blob.user_project
+
+ # When a Customer Managed Encryption Key is used to encrypt Cloud Storage object
+ # at rest, object resource metadata will store the version of the Key Management
+ # Service cryptographic material. If a Blob instance with KMS Key metadata set is
+ # used to upload a new version of the object then the existing kmsKeyName version
+ # value can't be used in the upload request and the client instead ignores it.
+ if blob.kms_key_name is not None and "cryptoKeyVersions" not in blob.kms_key_name:
+ headers["x-goog-encryption-kms-key-name"] = blob.kms_key_name
+
+ container = XMLMPUContainer(url, filename, headers=headers)
+ container.initiate(transport=transport, content_type=content_type)
+ upload_id = container.upload_id
+
+ size = os.path.getsize(filename)
+ num_of_parts = -(size // -chunk_size) # Ceiling division
+
+ pool_class, needs_pickling = _get_pool_class_and_requirements(worker_type)
+ # Pickle the blob ahead of time (just once, not once per chunk) if needed.
+ maybe_pickled_client = _pickle_client(client) if needs_pickling else client
+
+ futures = []
+
+ with pool_class(max_workers=max_workers) as executor:
+
+ for part_number in range(1, num_of_parts + 1):
+ start = (part_number - 1) * chunk_size
+ end = min(part_number * chunk_size, size)
+
+ futures.append(
+ executor.submit(
+ _upload_part,
+ maybe_pickled_client,
+ url,
+ upload_id,
+ filename,
+ start=start,
+ end=end,
+ part_number=part_number,
+ checksum=checksum,
+ headers=headers,
+ )
+ )
+
+ concurrent.futures.wait(
+ futures, timeout=deadline, return_when=concurrent.futures.ALL_COMPLETED
+ )
+
+ try:
+ # Harvest results and raise exceptions.
+ for future in futures:
+ part_number, etag = future.result()
+ container.register_part(part_number, etag)
+
+ container.finalize(blob._get_transport(client))
+ except Exception:
+ container.cancel(blob._get_transport(client))
+ raise
+
+
+def _upload_part(
+ maybe_pickled_client,
+ url,
+ upload_id,
+ filename,
+ start,
+ end,
+ part_number,
+ checksum,
+ headers,
+):
+ """Helper function that runs inside a thread or subprocess to upload a part.
+
+ `maybe_pickled_client` is either a Client (for threads) or a specially
+ pickled Client (for processes) because the default pickling mangles Client
+ objects."""
+
+ if isinstance(maybe_pickled_client, Client):
+ client = maybe_pickled_client
+ else:
+ client = pickle.loads(maybe_pickled_client)
+ part = XMLMPUPart(
+ url,
+ upload_id,
+ filename,
+ start=start,
+ end=end,
+ part_number=part_number,
+ checksum=checksum,
+ headers=headers,
+ )
+ part.upload(client._http)
+ return (part_number, part.etag)
+
+
+def _headers_from_metadata(metadata):
+ """Helper function to translate object metadata into a header dictionary."""
+
+ headers = {}
+ # Handle standard writable metadata
+ for key, value in metadata.items():
+ if key in METADATA_HEADER_TRANSLATION:
+ headers[METADATA_HEADER_TRANSLATION[key]] = value
+ # Handle custom metadata
+ if "metadata" in metadata:
+ for key, value in metadata["metadata"].items():
+ headers["x-goog-meta-" + key] = value
+ return headers
+
+
 def _download_and_write_chunk_in_place(
  maybe_pickled_blob, filename, start, end, download_kwargs
 ):
+ """Helper function that runs inside a thread or subprocess.
+
+ `maybe_pickled_blob` is either a Blob (for threads) or a specially pickled
+ Blob (for processes) because the default pickling mangles Client objects
+ which are attached to Blobs."""
+
  if isinstance(maybe_pickled_blob, Blob):
  blob = maybe_pickled_blob
  else:
@@ -863,9 +1121,9 @@ def _call_method_on_maybe_pickled_blob(
 ):
  """Helper function that runs inside a thread or subprocess.
 
- `maybe_pickled_blob` is either a blob (for threads) or a specially pickled
- blob (for processes) because the default pickling mangles clients which are
- attached to blobs."""
+ `maybe_pickled_blob` is either a Blob (for threads) or a specially pickled
+ Blob (for processes) because the default pickling mangles Client objects
+ which are attached to Blobs."""
 
  if isinstance(maybe_pickled_blob, Blob):
  blob = maybe_pickled_blob
@@ -894,8 +1152,8 @@ def _reduce_client(cl):
  )
 
 
-def _pickle_blob(blob):
- """Pickle a Blob (and its Bucket and Client) and return a bytestring."""
+def _pickle_client(obj):
+ """Pickle a Client or an object that owns a Client (like a Blob)"""
 
  # We need a custom pickler to process Client objects, which are attached to
  # Buckets (and therefore to Blobs in turn). Unfortunately, the Python
@@ -907,7 +1165,7 @@ def _pickle_blob(blob):
  p = pickle.Pickler(f)
  p.dispatch_table = copyreg.dispatch_table.copy()
  p.dispatch_table[Client] = _reduce_client
- p.dump(blob)
+ p.dump(obj)
  return f.getvalue()
 
 

setup.py

@@ -31,7 +31,7 @@
  "google-auth >= 1.25.0, < 3.0dev",
  "google-api-core >= 1.31.5, <3.0.0dev,!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0",
  "google-cloud-core >= 2.3.0, < 3.0dev",
- "google-resumable-media >= 2.3.2",
+ "google-resumable-media >= 2.6.0",
  "requests >= 2.18.0, < 3.0.0dev",
 ]
 extras = {"protobuf": ["protobuf<5.0.0dev"]}