patch for ci test

dstaay-fb · facebook-github-bot · commit c4021f6f8152 · 2025-07-29T20:38:42.000-07:00
Differential Revision: D79229308
diff --git a/cuda-sys/build.rs b/cuda-sys/build.rs
@@ -74,7 +74,7 @@ fn emit_cuda_link_directives(cuda_home: &str) {
 }
 
 fn python_env_dirs() -> (Option<String>, Option<String>) {
-    let output = std::process::Command::new(PathBuf::from("python"))
+    let output = std::process::Command::new(PathBuf::from("python3"))
         .arg("-c")
         .arg(PYTHON_PRINT_DIRS)
         .output()
@@ -94,13 +94,13 @@ fn python_env_dirs() -> (Option<String>, Option<String>) {
 }
 
 fn main() {
+    // Start building the bindgen configuration
     let mut builder = bindgen::Builder::default()
         // The input header we would like to generate bindings for
         .header("src/wrapper.h")
         .clang_arg("-x")
         .clang_arg("c++")
         .clang_arg("-std=gnu++20")
-        .clang_arg(format!("-I{}/include", find_cuda_home().unwrap()))
         .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
         // Allow the specified functions and types
         .allowlist_function("cu.*")
@@ -113,6 +113,21 @@ fn main() {
             is_global: false,
         });
 
+    // Add CUDA include path if available
+    if let Some(cuda_home) = find_cuda_home() {
+        let cuda_include_path = format!("{}/include", cuda_home);
+        if Path::new(&cuda_include_path).exists() {
+            builder = builder.clang_arg(format!("-I{}", cuda_include_path));
+        } else {
+            eprintln!(
+                "Warning: CUDA include directory not found at {}",
+                cuda_include_path
+            );
+        }
+    } else {
+        eprintln!("Warning: CUDA home directory not found. Continuing without CUDA include path.");
+    }
+
     // Include headers and libs from the active environment.
     let (include_dir, lib_dir) = python_env_dirs();
     if let Some(include_dir) = include_dir {
@@ -129,15 +144,28 @@ fn main() {
     }
 
     // Write the bindings to the $OUT_DIR/bindings.rs file
-    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
-    builder
-        .generate()
-        .expect("Unable to generate bindings")
-        .write_to_file(out_path.join("bindings.rs"))
-        .expect("Couldn't write bindings!");
-
-    println!("cargo:rustc-link-lib=cuda");
-    println!("cargo:rustc-link-lib=cudart");
-    println!("cargo::rustc-cfg=cargo");
-    println!("cargo::rustc-check-cfg=cfg(cargo)");
+    match env::var("OUT_DIR") {
+        Ok(out_dir) => {
+            let out_path = PathBuf::from(out_dir);
+            match builder.generate() {
+                Ok(bindings) => match bindings.write_to_file(out_path.join("bindings.rs")) {
+                    Ok(_) => {
+                        println!("cargo:rustc-link-lib=cuda");
+                        println!("cargo:rustc-link-lib=cudart");
+                        println!("cargo::rustc-cfg=cargo");
+                        println!("cargo::rustc-check-cfg=cfg(cargo)");
+                    }
+                    Err(e) => eprintln!("Warning: Couldn't write bindings: {}", e),
+                },
+                Err(e) => eprintln!("Warning: Unable to generate bindings: {}", e),
+            }
+        }
+        Err(_) => {
+            // When running via buck2 run, OUT_DIR might not be set
+            // This is expected and not an error - we're just running the script directly
+            // The actual build will happen later with cargo
+            println!("cargo:rustc-link-lib=cuda");
+            println!("cargo:rustc-link-lib=cudart");
+        }
+    }
 }
diff --git a/rdmaxcel-sys/Cargo.toml b/rdmaxcel-sys/Cargo.toml
@@ -4,7 +4,7 @@ version = "0.0.1"
 authors = ["Facebook"]
 edition = "2021"
 license = "MIT"
-links = "ibverbs"
+links = ["ibverbs", "cuda"]
 
 [dependencies]
 cxx = "1.0.119"
diff --git a/rdmaxcel-sys/build.rs b/rdmaxcel-sys/build.rs
@@ -13,6 +13,12 @@ use std::path::PathBuf;
 use glob::glob;
 use which::which;
 
+const PYTHON_PRINT_DIRS: &str = r"
+import sysconfig
+print('PYTHON_INCLUDE_DIR:', sysconfig.get_config_var('INCLUDEDIR'))
+print('PYTHON_LIB_DIR:', sysconfig.get_config_var('LIBDIR'))
+";
+
 // Translated from torch/utils/cpp_extension.py
 fn find_cuda_home() -> Option<String> {
     // Guess #1
@@ -67,6 +73,26 @@ fn emit_cuda_link_directives(cuda_home: &str) {
     println!("cargo:rustc-link-lib=cudart");
 }
 
+fn python_env_dirs() -> (Option<String>, Option<String>) {
+    let output = std::process::Command::new(PathBuf::from("python3"))
+        .arg("-c")
+        .arg(PYTHON_PRINT_DIRS)
+        .output()
+        .unwrap_or_else(|_| panic!("error running python"));
+
+    let mut include_dir = None;
+    let mut lib_dir = None;
+    for line in String::from_utf8_lossy(&output.stdout).lines() {
+        if let Some(path) = line.strip_prefix("PYTHON_INCLUDE_DIR: ") {
+            include_dir = Some(path.to_string());
+        }
+        if let Some(path) = line.strip_prefix("PYTHON_LIB_DIR: ") {
+            lib_dir = Some(path.to_string());
+        }
+    }
+    (include_dir, lib_dir)
+}
+
 fn main() {
     // Tell cargo to look for shared libraries in the specified directory
     println!("cargo:rustc-link-search=/usr/lib");
@@ -78,30 +104,65 @@ fn main() {
     // Link against the mlx5 library
     println!("cargo:rustc-link-lib=mlx5");
 
-    // Link against cuda library library
-    if let Some(cuda_home) = find_cuda_home() {
-        emit_cuda_link_directives(&cuda_home);
-    }
-
     // Tell cargo to invalidate the built crate whenever the wrapper changes
     println!("cargo:rerun-if-changed=src/rdmaxcel.h");
 
-    // Add cargo metadata
-    println!("cargo:rustc-cfg=cargo");
-    println!("cargo:rustc-check-cfg=cfg(cargo)");
+    // Get the directory of the current crate
+    let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| {
+        // For buck2 run, we know the package is in fbcode/monarch/rdmaxcel-sys
+        // Get the fbsource directory from the current directory path
+        let current_dir = std::env::current_dir().expect("Failed to get current directory");
+        let current_path = current_dir.to_string_lossy();
+
+        // Find the fbsource part of the path
+        if let Some(fbsource_pos) = current_path.find("fbsource") {
+            let fbsource_path = &current_path[..fbsource_pos + "fbsource".len()];
+            format!("{}/fbcode/monarch/rdmaxcel-sys", fbsource_path)
+        } else {
+            // If we can't find fbsource in the path, just use the current directory
+            format!("{}/src", current_dir.to_string_lossy())
+        }
+    });
 
-    // The bindgen::Builder is the main entry point to bindgen
-    let bindings = bindgen::Builder::default()
+    // Create the absolute path to the header file
+    let header_path = format!("{}/src/rdmaxcel.h", manifest_dir);
+
+    // Check if the header file exists
+    if !Path::new(&header_path).exists() {
+        panic!("Header file not found at {}", header_path);
+    }
+
+    // Start building the bindgen configuration
+    let mut builder = bindgen::Builder::default()
         // The input header we would like to generate bindings for
-        .header("src/rdmaxcel.h")
+        .header(&header_path)
+        .clang_arg("-x")
+        .clang_arg("c++")
+        .clang_arg("-std=gnu++20")
+        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
         // Allow the specified functions, types, and variables
         .allowlist_function("ibv_.*")
         .allowlist_function("mlx5dv_.*")
         .allowlist_function("mlx5_wqe_.*")
+        .allowlist_function("create_qp")
+        .allowlist_function("create_mlx5dv_.*")
+        .allowlist_function("register_cuda_memory")
+        .allowlist_function("db_ring")
+        .allowlist_function("cqe_poll")
+        .allowlist_function("send_wqe")
+        .allowlist_function("recv_wqe")
+        .allowlist_function("launch_db_ring")
+        .allowlist_function("launch_cqe_poll")
+        .allowlist_function("launch_send_wqe")
+        .allowlist_function("launch_recv_wqe")
         .allowlist_type("ibv_.*")
         .allowlist_type("mlx5dv_.*")
         .allowlist_type("mlx5_wqe_.*")
+        .allowlist_type("cqe_poll_result_t")
+        .allowlist_type("wqe_params_t")
+        .allowlist_type("cqe_poll_params_t")
         .allowlist_var("MLX5_.*")
+        .allowlist_var("IBV_.*")
         // Block specific types that are manually defined in lib.rs
         .blocklist_type("ibv_wc")
         .blocklist_type("mlx5_wqe_ctrl_seg")
@@ -118,15 +179,58 @@ fn main() {
         .constified_enum_module("ibv_wr_opcode")
         .constified_enum_module("ibv_wc_status")
         .derive_default(true)
-        .prepend_enum_name(false)
-        // Finish the builder and generate the bindings
-        .generate()
-        // Unwrap the Result and panic on failure
-        .expect("Unable to generate bindings");
+        .prepend_enum_name(false);
+
+    // Add CUDA include path if available
+    if let Some(cuda_home) = find_cuda_home() {
+        let cuda_include_path = format!("{}/include", cuda_home);
+        if Path::new(&cuda_include_path).exists() {
+            builder = builder.clang_arg(format!("-I{}", cuda_include_path));
+        } else {
+            eprintln!(
+                "Warning: CUDA include directory not found at {}",
+                cuda_include_path
+            );
+        }
+    } else {
+        eprintln!("Warning: CUDA home directory not found. Continuing without CUDA include path.");
+    }
+
+    // Include headers and libs from the active environment.
+    let (include_dir, lib_dir) = python_env_dirs();
+    if let Some(include_dir) = include_dir {
+        builder = builder.clang_arg(format!("-I{}", include_dir));
+    }
+    if let Some(lib_dir) = lib_dir {
+        println!("cargo::rustc-link-search=native={}", lib_dir);
+        // Set cargo metadata to inform dependent binaries about how to set their
+        // RPATH (see controller/build.rs for an example).
+        println!("cargo::metadata=LIB_PATH={}", lib_dir);
+    }
+    if let Some(cuda_home) = find_cuda_home() {
+        emit_cuda_link_directives(&cuda_home);
+    }
+
+    // Generate bindings
+    let bindings = builder.generate().expect("Unable to generate bindings");
 
     // Write the bindings to the $OUT_DIR/bindings.rs file
-    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
-    bindings
-        .write_to_file(out_path.join("bindings.rs"))
-        .expect("Couldn't write bindings!");
+    match env::var("OUT_DIR") {
+        Ok(out_dir) => {
+            let out_path = PathBuf::from(out_dir);
+            match bindings.write_to_file(out_path.join("bindings.rs")) {
+                Ok(_) => {
+                    println!("cargo:rustc-cfg=cargo");
+                    println!("cargo:rustc-check-cfg=cfg(cargo)");
+                }
+                Err(e) => eprintln!("Warning: Couldn't write bindings: {}", e),
+            }
+        }
+        Err(_) => {
+            // When running via buck2 run, OUT_DIR might not be set
+            // This is expected and not an error - we're just running the script directly
+            // The actual build will happen later with cargo
+            println!("Note: OUT_DIR not set, skipping bindings file generation");
+        }
+    }
 }
diff --git a/rdmaxcel-sys/src/rdmaxcel.h b/rdmaxcel-sys/src/rdmaxcel.h
@@ -9,7 +9,24 @@
 #ifndef RDMAXCEL_H
 #define RDMAXCEL_H
 
+// Conditionally include CUDA headers or provide stub definitions
+#if defined(HAVE_CUDA) || defined(__CUDACC__)
 #include <cuda_runtime.h>
+#else
+// Stub definitions for CUDA types when CUDA is not available
+#ifdef __cplusplus
+extern "C" {
+#endif
+typedef int cudaError_t;
+typedef void* CUdeviceptr;
+#define __host__
+#define __device__
+#define __global__
+#ifdef __cplusplus
+}
+#endif
+#endif
+
 #include <infiniband/mlx5dv.h>
 #include <infiniband/verbs.h>
 
