Skip to content

Commit

Permalink
feat(serve): Add graceful Ctrl+C handling with exit message (#1238)
Browse files Browse the repository at this point in the history
# Graceful Ctrl+C Handling

Implements reliable Ctrl+C handling for BAML's Python interface with a
clean shutdown message. This fixes issues with signal handling across
the Python/Rust boundary and improves user experience when stopping BAML
processes.

## Changes
- Added custom signal handling using the `ctrlc` crate
- Implemented clean "Shutting Down BAML..." message on interrupt
- Bypassed Python's signal handling to ensure reliable behavior
- Added proper exit code (130) following Unix conventions
- Added detailed documentation explaining the implementation

## Testing
- Tested Python CLI server start/stop
- Manual Ctrl+C testing in various scenarios
- Verified clean shutdown message display
- Tested cross-platform signal handling behavior

## Technical Details
- Uses `ctrlc` crate for reliable cross-platform signal handling
- Implements channel-based communication between signal handler and main
thread
- Follows Unix conventions for interrupt handling (exit code 130)
- Properly documented implementation rationale and technical decisions

## Notes
- While this is a workaround, it provides reliable interrupt handling
without requiring major architectural changes
- The implementation bypasses Python's signal handling entirely to avoid
conflicts
- Future improvements could involve deeper integration with BAML's
architecture

<!-- ELLIPSIS_HIDDEN -->

----

> [!IMPORTANT]
> Enhances BAML's Python interface with reliable Ctrl+C handling and
clean shutdown message using the `ctrlc` crate.
> 
>   - **Behavior**:
> - Implements custom signal handling to bypass Python's signal handlers
>     - Displays "Shutting Down BAML..." message on interrupt
>     - Ensures proper process termination with standard exit codes
>   - **Implementation**:
>     - Uses `ctrlc` crate for cross-platform signal handling
>     - Implements channel-based communication for safe signal handling
>     - Follows Unix conventions with exit code 130 for SIGINT
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for the latest commit. It will automatically update as commits are
pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->
  • Loading branch information
afyef authored Dec 23, 2024
1 parent 9170b89 commit 83e68f2
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 0 deletions.
57 changes: 57 additions & 0 deletions engine/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 3 additions & 0 deletions engine/language_client_python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,9 @@ internal-baml-codegen.workspace = true
env_logger.workspace = true
futures.workspace = true
indexmap.workspace = true
libc = "0.2"
log.workspace = true
ctrlc = "3.4"
# Consult https://pyo3.rs/main/migration for migration instructions
pyo3 = { version = "0.23.3", default-features = false, features = [
"abi3-py38",
Expand All @@ -44,6 +46,7 @@ regex.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio = { version = "1", features = ["full"] }
tokio-util = { version = "0.7", features = ["full"] }
tracing-subscriber = { version = "0.3.18", features = [
"json",
"env-filter",
Expand Down
46 changes: 46 additions & 0 deletions engine/language_client_python/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,55 @@ use pyo3::prelude::{pyfunction, pymodule, PyAnyMethods, PyModule, PyResult};
use pyo3::types::PyModuleMethods;
use pyo3::{wrap_pyfunction, Bound, Python};
use tracing_subscriber::{self, EnvFilter};
use ctrlc;

#[pyfunction]
fn invoke_runtime_cli(py: Python) -> PyResult<()> {
// SIGINT (Ctrl+C) Handling Implementation, an approach from @revidious
//
// Background:
// When running BAML through Python, we face a challenge where Python's default SIGINT handling
// can interfere with graceful shutdown. This is because:
// 1. Python has its own signal handlers that may conflict with Rust's
// 2. The PyO3 runtime can sometimes mask or delay interrupt signals
// 3. We need to ensure clean shutdown across the Python/Rust boundary
//
// Solution:
// We implement a custom signal handling mechanism using Rust's ctrlc crate that:
// 1. Bypasses Python's signal handling entirely
// 2. Provides consistent behavior across platforms
// 3. Ensures graceful shutdown with proper exit codes
// Note: While eliminating the root cause of SIGINT handling conflicts would be ideal,
// the source appears to be deeply embedded in BAML's architecture and PyO3's runtime.
// A proper fix would require extensive changes to how BAML handles signals across the
// Python/Rust boundary. For now, this workaround provides reliable interrupt handling
// without requiring major architectural changes but welp, this is a hacky solution.

// Create a channel for communicating between the signal handler and main thread
// This is necessary because signal handlers run in a separate context and
// need a safe way to communicate with the main program
let (interrupt_send, interrupt_recv) = std::sync::mpsc::channel();

// Install our custom Ctrl+C handler
// This will run in a separate thread when SIGINT is received
ctrlc::set_handler(move || {
println!("\nShutting Down BAML...");
// Notify the main thread through the channel
// Using ok() to ignore send errors if the receiver is already dropped
interrupt_send.send(()).ok();
}).expect("Error setting Ctrl-C handler");

// Monitor for interrupt signals in a separate thread
// This is necessary because we can't directly exit from the signal handler.

std::thread::spawn(move || {
if interrupt_recv.recv().is_ok() {
// Exit with code 130 (128 + SIGINT's signal number 2)
// This is the standard Unix convention for processes terminated by SIGINT
std::process::exit(130);
}
});

baml_cli::run_cli(
py.import("sys")?
.getattr("argv")?
Expand Down

0 comments on commit 83e68f2

Please sign in to comment.