fix ImportError: libpython3.10.so.1.0: cannot open shared object file: No such file or directory (#1616)

Summary:
ImportError: libpython3.10.so.1.0: cannot open shared object file: No such file or directory

to fix, we embeds RPATH entries into the compiled Rust shared libraries, telling the dynamic linker where to search for libpython3.10.so.1.0 at runtime (shown by D84991141 by shayne-fletcher)

in order to built  torchsys e2e (and handle both cuda 12.6 and 12.9), improve cuda path discovery

also fix bug in device_mesh

Pull Request resolved: #1616

Test Plan:
pip install torch
pip install "torchmonarch[examples]==0.1.0rc7"
(monarch) [[email protected] ~/monarch (test-import)]$ python docs/source/examples/distributed_tensors.py

Imported from GitHub, without a `Test Plan:` line.

Reviewed By: shayne-fletcher

Differential Revision: D85020138

Pulled By: colin2328

fbshipit-source-id: c701c762ea071cba513caf8d24f8f76b899ddbe8

Author: colin2328

build_utils/src/lib.rs

@@ -103,14 +103,14 @@ pub fn get_env_var_with_rerun(name: &str) -> Result<String, std::env::VarError>
 ///
 /// This function attempts to locate CUDA installation through:
 /// 1. CUDA_HOME environment variable
-/// 2. CUDA_PATH environment variable  
+/// 2. CUDA_PATH environment variable
 /// 3. Finding nvcc in PATH and deriving cuda home
 /// 4. Platform-specific default locations
 pub fn find_cuda_home() -> Option<String> {
     // Guess #1: Environment variables
-    let mut cuda_home = env::var("CUDA_HOME")
+    let mut cuda_home = get_env_var_with_rerun("CUDA_HOME")
         .ok()
-        .or_else(|| env::var("CUDA_PATH").ok());
+        .or_else(|| get_env_var_with_rerun("CUDA_PATH").ok());
 
     if cuda_home.is_none() {
         // Guess #2: Find nvcc in PATH
@@ -151,14 +151,18 @@ pub fn discover_cuda_config() -> Result<CudaConfig, BuildError> {
         lib_dirs: Vec::new(),
     };
 
-    // Add standard include directory
-    let include_dir = cuda_home_path.join("include");
-    if include_dir.exists() {
-        config.include_dirs.push(include_dir);
+    // Add standard include directories
+    // Check both old-style (include) and new-style (targets/x86_64-linux/include) CUDA installations
+    for include_subdir in &["include", "targets/x86_64-linux/include"] {
+        let include_dir = cuda_home_path.join(include_subdir);
+        if include_dir.exists() {
+            config.include_dirs.push(include_dir);
+        }
     }
 
     // Add standard library directories
-    for lib_subdir in &["lib64", "lib", "lib/x64"] {
+    // Check both old-style (lib64, lib) and new-style (targets/x86_64-linux/lib) CUDA installations
+    for lib_subdir in &["lib64", "lib", "lib/x64", "targets/x86_64-linux/lib"] {
         let lib_dir = cuda_home_path.join(lib_subdir);
         if lib_dir.exists() {
             config.lib_dirs.push(lib_dir);
@@ -197,13 +201,12 @@ pub fn get_cuda_lib_dir() -> Result<String, BuildError> {
     // Try to deduce from CUDA configuration
     let cuda_config = discover_cuda_config()?;
     if let Some(cuda_home) = cuda_config.cuda_home {
-        let lib64_path = cuda_home.join("lib64");
-        if lib64_path.exists() {
-            return Ok(lib64_path.to_string_lossy().to_string());
-        }
-        let lib_path = cuda_home.join("lib");
-        if lib_path.exists() {
-            return Ok(lib_path.to_string_lossy().to_string());
+        // Check both old-style and new-style CUDA library paths
+        for lib_subdir in &["lib64", "lib", "targets/x86_64-linux/lib"] {
+            let lib_path = cuda_home.join(lib_subdir);
+            if lib_path.exists() {
+                return Ok(lib_path.to_string_lossy().to_string());
+            }
         }
     }
 

python/monarch/_src/actor/v1/proc_mesh.py

@@ -292,7 +292,7 @@ def _spawn_nonblocking(
         return self._spawn_nonblocking_on(self._proc_mesh, name, Class, *args, **kwargs)
 
     def to_table(self) -> str:
-        return self._maybe_device_mesh.to_table()
+        return self._device_mesh.to_table()
 
     def _spawn_nonblocking_on(
         self,

setup.py

@@ -124,6 +124,7 @@ def run(self):
     # Always include the active env's lib (Conda-safe)
     conda_lib = os.path.join(sys.prefix, "lib")
 
+    # Only use LIBDIR if it actually contains the current libpython
     ldlib = sysconfig.get_config_var("LDLIBRARY") or ""
     libdir = sysconfig.get_config_var("LIBDIR") or ""
     py_lib = ""

torch-sys-cuda/build.rs

@@ -28,7 +28,9 @@ fn main() {
     let mut libtorch_lib_dir: Option<PathBuf> = None;
     let mut cxx11_abi = None;
     let mut cuda_home: Option<PathBuf> = None;
-    let python_interpreter = PathBuf::from("python");
+    let python_interpreter = std::env::var("PYO3_PYTHON")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| PathBuf::from("python"));
 
     let use_pytorch_apis = build_utils::get_env_var_with_rerun("TORCH_SYS_USE_PYTORCH_APIS")
         .unwrap_or_else(|_| "1".to_owned());

torch-sys/build.rs

@@ -26,7 +26,9 @@ fn main() {
     let mut libtorch_include_dirs: Vec<PathBuf> = vec![];
     let mut libtorch_lib_dir: Option<PathBuf> = None;
     let mut cxx11_abi = None;
-    let python_interpreter = PathBuf::from("python");
+    let python_interpreter = std::env::var("PYO3_PYTHON")
+        .map(PathBuf::from)
+        .unwrap_or_else(|_| PathBuf::from("python"));
 
     let use_pytorch_apis = build_utils::get_env_var_with_rerun("TORCH_SYS_USE_PYTORCH_APIS")
         .unwrap_or_else(|_| "1".to_owned());