diff --git a/python/probing/cli/__main__.py b/python/probing/cli/__main__.py index 19674a6..60402aa 100644 --- a/python/probing/cli/__main__.py +++ b/python/probing/cli/__main__.py @@ -1,19 +1,14 @@ +import os import sys +# Set CLI mode before any probing imports to skip probe initialization +os.environ["PROBING_CLI_MODE"] = "1" + def main(): """Entry point for the probing CLI command.""" - # Import probing module - this will load _core via maturin - # cli_main is exported from probing.__init__.py which imports it from _core import probing - # cli_main is available from probing module (exported from __init__.py) - if not hasattr(probing, "cli_main"): - raise ImportError( - "cli_main is not available in probing module. " - "Please ensure the probing package is properly installed." - ) - probing.cli_main(["probing"] + sys.argv[1:]) diff --git a/python/probing_hook.py b/python/probing_hook.py index fe19c73..c15b465 100644 --- a/python/probing_hook.py +++ b/python/probing_hook.py @@ -115,10 +115,31 @@ def init_probing(): # In case of unexpected errors, don't enable probing +def is_probing_cli(): + """Check if the current process is the probing CLI itself.""" + if current_script == "probing": + os.environ["PROBING_CLI_MODE"] = "1" + return True + # Check for 'python -m probing.cli' + try: + import __main__ + + if hasattr(__main__, "__file__") and __main__.__file__: + if "probing" in __main__.__file__ and "cli" in __main__.__file__: + os.environ["PROBING_CLI_MODE"] = "1" + return True + except Exception: + pass + return False + + try: import re - if re.search("torchrun", current_script) is None: + # Skip initialization for: + # 1. torchrun (launcher script) + # 2. probing CLI itself (should not inject probes into itself) + if re.search("torchrun", current_script) is None and not is_probing_cli(): init_probing() except Exception as e: print(f"Error in probing hook: {e}", file=sys.stderr) diff --git a/src/lib.rs b/src/lib.rs index edb8ede..25a3886 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,7 +3,6 @@ extern crate ctor; use anyhow::Result; use pyo3::prelude::*; -use std::net::ToSocketAddrs; use probing_python::extensions::python::ExternalTable; use probing_python::features::config; @@ -145,8 +144,16 @@ fn setup_env_settings() { } } +const ENV_PROBING_CLI_MODE: &str = "PROBING_CLI_MODE"; + #[ctor] fn setup() { + // Skip initialization if running in CLI mode (e.g., probing ls) + // CLI commands should not inject probes into themselves + if std::env::var(ENV_PROBING_CLI_MODE).is_ok() { + return; + } + let pid = std::process::id(); eprintln!("Initializing probing module for process {pid} ...",); @@ -164,6 +171,11 @@ fn setup() { #[dtor] fn cleanup() { + // Skip cleanup if running in CLI mode (no probes were initialized) + if std::env::var(ENV_PROBING_CLI_MODE).is_ok() { + return; + } + if let Err(e) = probing_server::cleanup() { log::error!("Failed to cleanup unix socket: {e}"); }