fix: init mp pool & grpc client once, use os.sched_setaffinity

[chandra-siri]

fix: init mp pool & grpc client once, use os.sched_setaffinity

• mp pool is initalize once per benchmark run, otherwise time to init pool is affecting throughput calculations.
• grpc client should be intialized once per process in an event loop, otherwise processes get stuck .
• os.sched_setaffinity helps in pinning the process to a set of cores to avoid hard irq from nic.

[gemini-code-assist]

Summary of Changes

Hello @chandra-siri, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces significant performance enhancements and stability improvements to the microbenchmarks by optimizing resource initialization and process management. It ensures that critical components like multiprocessing pools and gRPC clients are set up efficiently, and worker processes are assigned dedicated CPU resources to minimize external interference, leading to more accurate and reliable benchmark results.

Highlights

• Performance Optimization: The multiprocessing pool is now initialized only once per benchmark run, preventing repeated setup overhead from affecting throughput calculations.
• Client Initialization: gRPC and Google Cloud Storage clients are now initialized once per worker process, resolving issues where processes would get stuck due to repeated client creation within event loops.
• CPU Affinity: Worker processes are now pinned to a specific range of CPU cores (20-189) using os.sched_setaffinity, which helps mitigate performance interference from hard interrupts from the network interface card (NIC).

Changelog

• tests/perf/microbenchmarks/writes/test_writes.py
  • Introduced global variables for asyncio event loops and gRPC/JSON clients to enable per-process, one-time initialization.
  • Added a new _worker_init function to handle the one-time setup of event loops, clients, and CPU affinity for each worker process.
  • Modified the _upload_files_worker function to utilize the globally initialized clients, removing redundant client creation within each worker call.
  • Updated the upload_files_mp_mc_wrapper function to accept a pre-initialized multiprocessing pool, decoupling pool creation from the wrapper's execution.
  • Refactored the main benchmark execution logic to create the multiprocessing pool once with an initializer, ensuring proper setup and teardown of worker resources.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors the multiprocessing setup in the write benchmarks to initialize the process pool and clients only once per run, optimizing benchmark results by preventing setup overhead. From a security standpoint, the analysis could not be fully completed as several files listed in SECURITY_ANALYSIS_TODO.md were unreadable, meaning no vulnerabilities were found in the accessible files, but a comprehensive review was hindered. For code improvements, consider making the os.sched_setaffinity CPU affinity setting more robust and portable, and ensuring graceful cleanup of resources in worker processes.

[gemini-code-assist]

The hardcoded core range range(20, 190) for os.sched_setaffinity is not portable and will raise a ValueError on systems with fewer than 190 cores, causing the benchmark to crash. Please consider making this more robust by checking against the number of available cores, for example by using os.cpu_count().

[gemini-code-assist]

The _worker_init function initializes an event loop and clients for each worker process, but there's no corresponding cleanup logic. This can lead to resource leaks. While process termination will clean up resources, a graceful shutdown is better practice. Consider using atexit.register() in _worker_init to call a cleanup function that closes the clients and the event loop before the worker process exits.