Skip to content

Commit

Permalink
feat(python): implement graceful ctrl+c handling
Browse files Browse the repository at this point in the history
This implements a custom signal handling mechanism using the ctrlc crate to: 1. Bypass Python's signal handling 2. Provide consistent behavior across platforms 3. Ensure graceful shutdown with proper exit codes

While this is a workaround, it provides reliable interrupt handling without requiring major architectural changes to BAML's signal handling.
  • Loading branch information
afyef committed Dec 14, 2024
1 parent 9f0b4d2 commit d1952d6
Showing 1 changed file with 46 additions and 0 deletions.
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 d1952d6

Please sign in to comment.