Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(serve): Add graceful Ctrl+C handling with exit message #1238

Open
wants to merge 4 commits into
base: canary
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading