Add benchmarks for downloading and uploading large objects, and improve CRC32 performance measurement

Author: chandra-siri

benchmarks/download_zb_m_proc_n_coro.py

@@ -0,0 +1,128 @@
+import argparse
+import asyncio
+from datetime import datetime, timedelta
+from io import BytesIO
+import os
+import time
+import threading
+import logging
+from concurrent.futures import ProcessPoolExecutor
+import multiprocessing
+
+from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
+from google.cloud.storage._experimental.asyncio.async_multi_range_downloader import (
+    AsyncMultiRangeDownloader,
+)
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from tests.perf.microbenchmarks.resource_monitor import ResourceMonitor
+
+async def download_one_async(bucket_name, object_name, download_size, chunk_size):
+    """Downloads a single object of size `download_size`, in chunks of `chunk_size`"""
+    logging.debug(f"Downloading {object_name} of size {download_size} in chunks of {chunk_size} from {bucket_name} from process {os.getpid()} and thread {threading.get_ident()}")
+    client = AsyncGrpcClient().grpc_client
+
+    mrd = AsyncMultiRangeDownloader(client, bucket_name, object_name)
+    
+    start_time = time.perf_counter()
+    await mrd.open()
+    
+    offset = 0
+    output_buffer = BytesIO()
+    while offset < download_size:
+        bytes_to_download = min(chunk_size, download_size - offset)    
+        await mrd.download_ranges([(offset, bytes_to_download, output_buffer)])
+        offset += bytes_to_download
+
+    assert output_buffer.getbuffer().nbytes == download_size, f"downloaded size incorrect for {object_name}"
+    
+    await mrd.close()
+    end_time = time.perf_counter()
+    return end_time - start_time
+
+def run_coroutines(bucket_name, object_names, download_size, chunk_size):
+    """Runs a number of coroutines to download objects concurrently and returns their latencies."""
+    async def main():
+        tasks = [
+            download_one_async(bucket_name, object_name, download_size, chunk_size)
+            for object_name in object_names
+        ]
+        return await asyncio.gather(*tasks)
+    
+    return asyncio.run(main())
+
+def main():
+    multiprocessing.set_start_method("spawn")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--bucket_name", type=str, default='chandrasiri-rs')
+    parser.add_argument("--download_size_mib", type=int, default=1024)  # 1 GiB
+    parser.add_argument("--chunk_size_mib", type=int, default=64)  # 100 MiB
+    parser.add_argument("--count", type=int, default=4, help="The total number of objects to download.")
+    parser.add_argument("-m", "--num_processes", type=int, default=2, help="Number of worker processes.")
+    parser.add_argument("-n", "--num_coros", type=int, default=2, help="Number of worker coroutines.")
+    parser.add_argument("--start_object_num", type=int, default=0)
+    parser.add_argument(
+        "--run_for_minutes",
+        type=int,
+        default=0,
+        help="Number of minutes to run this process. If not provided, it runs for 1 iteration.",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', stream=sys.stdout)
+
+    iteration = 0
+    start_time = datetime.now()
+
+    def should_run():
+        if args.run_for_minutes > 0:
+            return datetime.now() - start_time < timedelta(minutes=args.run_for_minutes)
+        else:
+            return iteration == 0
+
+    while should_run():
+        logging.info(f"\n--- Iteration {iteration + 1} ---")
+        iteration += 1
+        total_start_time = time.perf_counter()
+
+        total_objects = args.count
+        all_object_names = [f"py-sdk-mb-mt-{i}" for i in range(args.start_object_num, args.start_object_num + total_objects)]
+
+        all_latencies = []
+        m = ResourceMonitor()
+
+        with m, ProcessPoolExecutor(max_workers=args.num_processes) as executor:
+            futures = []
+            for i in range(0, len(all_object_names), args.num_coros):
+                objects = all_object_names[i:i + args.num_coros]
+                future = executor.submit(
+                    run_coroutines, args.bucket_name, objects, args.download_size_mib * 1024 * 1024, args.chunk_size_mib * 1024 * 1024
+                )
+                futures.append(future)
+
+            
+        for future in futures:
+            all_latencies.extend(future.result())
+
+        total_end_time = time.perf_counter()
+        total_time_taken = total_end_time - total_start_time
+        total_downloaded_bytes = args.download_size_mib * 1024 * 1024 * total_objects
+        # m.throughput_mb_s = (total_downloaded_bytes / (1024*1024)) / total_time_taken if total_time_taken else 0
+
+        # Calculate throughput per-download and get the average
+
+        average_throughput = total_downloaded_bytes / total_time_taken / (1024 * 1024)  # MiB/s
+        logging.info("\n--- Aggregate Results ---")
+        logging.info(f"Total objects to download: {total_objects}")
+        logging.info(f"Total data to download: {total_downloaded_bytes / (1024*1024*1024):.2f} GiB")
+        logging.info(f"Total time taken: {total_time_taken:.2f} seconds")
+        logging.info(f"Average throughput: {average_throughput:.2f} MB/s")
+
+        logging.info("\n--- Resource Monitor Results ---")
+        logging.info(f"Max CPU: {m.max_cpu}")
+        logging.info(f"Max Memory: {m.max_mem}")
+        logging.info(f"Throughput (MB/s): {m.throughput_mb_s}")
+        logging.info(f"vCPUs: {m.vcpus}")
+
+if __name__ == "__main__":
+    main()

benchmarks/download_zb_n_workers.py

@@ -14,13 +14,14 @@
 
 async def download_one_async(bucket_name, object_name, download_size, chunk_size):
     """Downloads a single object of size `download_size`, in chunks of `chunk_size`"""
-    print(f"Downloading {object_name} of size {download_size} in chunks of {chunk_size} from {bucket_name} from process {os.getpid()} and thread {threading.get_ident()}")
-    # raise NotImplementedError("This function is not yet implemented.")
     client = AsyncGrpcClient().grpc_client
 
-    # log_peak_memory_usage()
     mrd = AsyncMultiRangeDownloader(client, bucket_name, object_name)
     await mrd.open()
+    # during benchmarking I've keyboard interrupted some uploads.
+    # So the requested download size may be larger than the actual object size.
+    download_size = min(download_size, mrd.persisted_size)
+    print(f"Downloading {object_name} of size {download_size} in chunks of {chunk_size} from {bucket_name} from process {os.getpid()} and thread {threading.get_ident()}")
 
     # download in chunks of `chunk_size`
     offset = 0
@@ -29,7 +30,6 @@ async def download_one_async(bucket_name, object_name, download_size, chunk_size
         bytes_to_download = min(chunk_size, download_size - offset)    
         await mrd.download_ranges([(offset, bytes_to_download, output_buffer)])
         offset += bytes_to_download
-    # await mrd.download_ranges([(offset, 0, output_buffer)])
     assert output_buffer.getbuffer().nbytes == download_size, f"downloaded size incorrect for {object_name}"
 
     await mrd.close()
@@ -41,10 +41,10 @@ def download_one_sync(bucket_name, object_name, download_size, chunk_size):
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--bucket_name", type=str, default='chandrasiri-rs')
+    parser.add_argument("--bucket_name", type=str, default='chandrasiri-benchmarks-zb')
     parser.add_argument("--download_size", type=int, default=1024 * 1024 * 1024)  # 1 GiB
     parser.add_argument("--chunk_size", type=int, default=64 * 1024 * 1024)  # 100 MiB
-    parser.add_argument("--count", type=int, default=100)
+    parser.add_argument("--count", type=int, default=4)
     parser.add_argument("--start_object_num", type=int, default=0)
     parser.add_argument("-n", "--num_workers", type=int, default=2, help="Number of worker threads or processes.")
     parser.add_argument("--executor", type=str, choices=['thread', 'process'], default='process', help="Executor to use: 'thread' for ThreadPoolExecutor, 'process' for ProcessPoolExecutor")

benchmarks/upload_large_objects.py

@@ -0,0 +1,107 @@
+import argparse
+import asyncio
+import os
+import time
+
+from google.cloud.storage._experimental.asyncio.async_grpc_client import AsyncGrpcClient
+from google.cloud.storage._experimental.asyncio.async_appendable_object_writer import (
+    AsyncAppendableObjectWriter,
+)
+from google.cloud import storage
+
+
+
+def delete_blob(bucket_name, blob_name):
+    """Deletes a blob from the bucket using the standard synchronous client."""
+    try:
+        storage_client = storage.Client()
+        bucket = storage_client.bucket(bucket_name)
+        blob = bucket.blob(blob_name)
+        print(f"--- Deleting blob {blob_name} from bucket {bucket_name}. ---", flush=True)
+        blob.delete()
+        print(f"--- Blob {blob_name} deleted. ---", flush=True)
+    except Exception as e:
+        print(f"--- Error deleting blob {blob_name}: {e} ---", flush=True)
+
+
+async def upload_one(client, bucket_name, object_name, upload_size, chunk_size):
+    """Uploads a single object of size `upload_size`, in chunks of `chunk_size`"""
+    print(
+        f"Uploading {object_name} of size {upload_size / (1024**3):.0f} GiB in chunks of {chunk_size / (1024**3):.0f} GiB to {bucket_name}",
+        flush=True
+    )
+
+    writer = AsyncAppendableObjectWriter(
+        client=client, bucket_name=bucket_name, object_name=object_name, writer_options={"FLUSH_INTERVAL_BYTES": chunk_size}
+    )
+
+    await writer.open()
+
+    start_time = time.perf_counter()
+
+    # Using a pre-allocated buffer for performance instead of os.urandom.
+    chunk = os.urandom(chunk_size)
+    uploaded_bytes = 0
+    while uploaded_bytes < upload_size:
+        # In a real-world scenario, you would read data from a source here.
+        # For this benchmark, we just send the same chunk repeatedly.
+        bytes_to_upload = min(chunk_size, upload_size - uploaded_bytes)
+        await writer.append(chunk[:bytes_to_upload])
+        uploaded_bytes += bytes_to_upload
+
+    response = await writer.close(finalize_on_close=True)
+    print(f"Upload response for {object_name}: {response}", flush=True)
+
+    end_time = time.perf_counter()
+    latency = end_time - start_time
+    throughput = (upload_size / latency) / (10**6)  # MB/s
+
+    print(f"Finished uploading {object_name}", flush=True)
+    print(f"Latency: {latency:.2f} seconds", flush=True)
+    print(f"Throughput: {throughput:.2f} MB/s", flush=True)
+
+
+async def main():
+    parser = argparse.ArgumentParser(
+        description="Benchmark large object uploads and deletions."
+    )
+    parser.add_argument(
+        "--bucket_name",
+        type=str,
+        default="chandrasiri-benchmarks-zb",
+        help="The GCS bucket to upload to.",
+    )
+    args = parser.parse_args()
+
+    async_client = AsyncGrpcClient().grpc_client
+    chunk_size_bytes = 1 * 1024 * 1024 * 1024  # 1 GiB
+
+    # Loop from 100 GiB to 1 TiB (1024 GiB) in 100 GiB increments.
+    for i in range(2, 5):
+        size_gib = i*100
+        upload_size_bytes = size_gib * 1024 * 1024 * 1024
+        object_name = f"large-upload-benchmark-{size_gib}gib"
+
+        print("\n" + "=" * 60, flush=True)
+        print(f"Starting Test Case: {size_gib} GiB Object Upload", flush=True)
+        print("=" * 60 + "\n", flush=True)
+
+        try:
+            await upload_one(
+                async_client,
+                args.bucket_name,
+                object_name,
+                upload_size_bytes,
+                chunk_size_bytes,
+            )
+        except Exception as e:
+            print(f"An error occurred during the upload for {object_name}: {e}", flush=True)
+        finally:
+            # Ensure the object is deleted after the test, even if the upload failed.
+            # We run the synchronous delete function in a separate thread to avoid
+            # blocking the asyncio event loop.
+            await asyncio.to_thread(delete_blob, args.bucket_name, object_name)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

crc32_bench.py

@@ -0,0 +1,47 @@
+import timeit
+import os
+import google_crc32c
+from google_crc32c import Checksum
+
+def benchmark():
+    # Setup: Create a 1MB random payload
+    data = os.urandom(2 * 1024 * 1024)
+    
+    # Method 1: The old approach (Object instantiation + digest + conversion)
+    def method_old():
+        # Note: Checksum(data).digest() returns bytes, which we convert to int
+        return int.from_bytes(Checksum(data).digest(), "big")
+
+    # Method 2: The new approach (Direct C-extension call)
+    def method_new():
+        return google_crc32c.value(data)
+
+    # Verify they produce the same result
+    assert method_old() == method_new(), "Checksums do not match!"
+
+    # Configuration for timeit
+    iterations = 1000
+    
+    print(f"Benchmarking with {len(data)/1024:.0f}KB payload over {iterations} iterations...")
+
+    # Measure Method 1
+    time_old = timeit.timeit(method_old, number=iterations)
+    print(f"Old method: {time_old:.4f} seconds total")
+
+    # Measure Method 2
+    time_new = timeit.timeit(method_new, number=iterations)
+    print(f"New method: {time_new:.4f} seconds total")
+
+    # Calculate difference
+    if time_new > 0:
+        speedup = time_old / time_new
+        print(f"\nResult: The new method is {speedup:.2f}x faster")
+        print(f"Difference per call: {(time_old - time_new) / iterations * 1000:.4f} ms")
+    else:
+        print("\nNew method was too fast to measure accurately.")
+
+if __name__ == "__main__":
+    try:
+        benchmark()
+    except ImportError:
+        print("Error: google-crc32c is not installed. Run 'pip install google-crc32c'")

google/cloud/storage/_experimental/asyncio/async_appendable_object_writer.py

@@ -293,8 +293,8 @@ async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object]
 
         if finalize_on_close:
             await self.finalize()
-        # else:
-        #     await self.flush()
+        else:
+            await self.flush()
 
         await self.write_obj_stream.close()
 

samples/snippets/storage_list_files_with_prefix.py

@@ -46,15 +46,6 @@ def list_blobs_with_prefix(bucket_name, prefix, delimiter=None):
     that lists the "subfolders" under `a/`:
 
         a/b/
-
-
-    Note: If you only want to list prefixes a/b/ and don't want to iterate over
-    blobs, you can do
-
-    ```
-    for page in blobs.pages:
-        print(page.prefixes)
-    ```
     """
 
     storage_client = storage.Client()
@@ -66,22 +57,16 @@ def list_blobs_with_prefix(bucket_name, prefix, delimiter=None):
 
     # Note: The call returns a response only when the iterator is consumed.
     print("Blobs:")
-    count = 0
     for blob in blobs:
-        if count % 500 == 0:
-            print(blob.name)
-        count += 1
-    print(f"Total blobs with prefix {prefix}: {count}")
-
-        # blob.delete()
+        print(blob.name)
 
-    # if delimiter:
-    #     print("Prefixes:")
-    #     for prefix in blobs.prefixes:
-    #         print(prefix)
+    if delimiter:
+        print("Prefixes:")
+        for prefix in blobs.prefixes:
+            print(prefix)
 
 
 # [END storage_list_files_with_prefix]
 
 if __name__ == "__main__":
-    list_blobs_with_prefix(bucket_name=sys.argv[1], prefix=sys.argv[2])
+    list_blobs_with_prefix(bucket_name=sys.argv[1], prefix=sys.argv[2])