Fix kernel timeout issues with pool management

[abhishek-anand]

Summary

• Implements kernel pool (2-5 kernels) to prevent single kernel bottlenecks
• Extends WebSocket timeouts to 10 minutes for long operations
• Adds retry logic with exponential backoff for failed executions
• Background health monitoring automatically replaces unresponsive kernels

Test plan

• Verify container starts with kernel pool initialization
• Test long-running operations (package installs, ML processing)
• Confirm automatic kernel recovery on failures
• Check concurrent execution handling

[gemini-code-assist]

Summary of Changes

Hello @abhishek-anand, 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 aims to significantly improve the robustness and performance of the Jupyter kernel management by introducing a kernel pooling mechanism. This change addresses potential single kernel bottlenecks and enhances the reliability of long-running operations by implementing retry logic with exponential backoff and extending WebSocket timeouts. Additionally, it ensures continuous availability through background health monitoring and automatic kernel replacement.

Highlights

• Kernel Pooling: Implemented a dynamic kernel pool that maintains between 2 and 5 active Jupyter kernels. This prevents single kernel bottlenecks and allows for concurrent execution of tasks. Kernels are acquired from the pool when needed and released back upon completion.
• Extended WebSocket Timeouts: Increased the WebSocket timeout for Jupyter kernel connections to 10 minutes (600 seconds). This change is crucial for accommodating long-running operations such as package installations or machine learning processing, preventing premature disconnections.
• Robust Retry Logic with Exponential Backoff: Introduced a execute_with_retry function that wraps kernel execution. This function retries failed operations up to 3 times, using an exponential backoff strategy (base 2) to space out retries, improving resilience against transient issues.
• Background Kernel Health Monitoring: A new background task continuously monitors the health of kernels in the pool. Unresponsive or failed kernels are automatically detected, removed, and replaced, ensuring a healthy and available pool of kernels.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 or fill out our survey to provide feedback.

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 introduces a robust kernel pooling mechanism to improve the stability and performance of Jupyter kernel interactions. The changes include a KernelPool class for managing kernel lifecycle, retry logic with exponential backoff for execution, and a background health checker to replace unresponsive kernels. The overall approach is solid and addresses the timeout issues described.

My review includes several suggestions to improve the implementation further:

• Removing unused code and imports to enhance clarity.
• Replacing blocking I/O calls with their asynchronous counterparts to prevent blocking the event loop.
• Using specific, custom exceptions instead of generic Exception for better error handling.
• Simplifying the logic for replacing failed kernels by centralizing it in the health check loop.

These changes will make the new kernel management system more maintainable and robust.