Enable RDMA tests in OSS CI (#1255)

Summary:
Pull Request resolved: #1255

- Ensure `cargo test -p monarch_rdma` runs on OSS CI
- Break out new testing infra into its own runner (see rationale below).
- `cargo test` in general will work for repo, but we have a lot of broken tests to resolve.

Given test suites getting relatively large and different paths (ie. cargo test vs pytest/pyright), which was actually helpful since some flakiness within python path having some long running examples.   Trade off is slightly more resources; so can revisit later if issue.

 {F1982168077}

Reviewed By: dulinriley

Differential Revision: D82699014

Author: dstaay-fb

.github/workflows/ci.yml

@@ -22,30 +22,38 @@ jobs:
     name: Build CPU
     uses: ./.github/workflows/build-cpu.yml
 
-  test-cuda:
-    name: Test CUDA
+  test-cpu-python:
+    name: Test CPU Python
+    needs: build-cpu
+    uses: ./.github/workflows/test-cpu-python.yml
+    with:
+      artifact-name: monarch-cpu-${{ github.sha }}
+
+  test-gpu-python:
+    name: Test GPU Python
     needs: build-cuda
-    uses: ./.github/workflows/test-cuda.yml
+    uses: ./.github/workflows/test-gpu-python.yml
     with:
       artifact-name: monarch-cuda-${{ github.sha }}
 
-  test-cpu:
-    name: Test CPU
-    needs: build-cpu
-    uses: ./.github/workflows/test-cpu.yml
+  test-gpu-rust:
+    name: Test GPU Rust
+    needs: build-cuda
+    uses: ./.github/workflows/test-gpu-rust.yml
     with:
-      artifact-name: monarch-cpu-${{ github.sha }}
+      artifact-name: monarch-cuda-${{ github.sha }}
 
   status-check:
     name: Status Check
     runs-on: ubuntu-latest
-    needs: [test-cuda, test-cpu]
+    needs: [test-cpu-python, test-gpu-python, test-gpu-rust]
     if: always()
     steps:
       - name: Check all jobs status
         run: |
-          if [[ "${{ needs.test-cuda.result }}" != "success" ]] ||
-             [[ "${{ needs.test-cpu.result }}" != "success" ]]; then
+          if [[ "${{ needs.test-cpu-python.result }}" != "success" ]] ||
+             [[ "${{ needs.test-gpu-python.result }}" != "success" ]] ||
+             [[ "${{ needs.test-gpu-rust.result }}" != "success" ]]; then
             echo "One or more jobs failed"
             exit 1
           else

.github/workflows/test-cpu.yml renamed to .github/workflows/test-cpu-python.yml

@@ -1,4 +1,4 @@
-name: Test CPU
+name: Test CPU Python
 
 on:
   workflow_call:
@@ -9,12 +9,12 @@ on:
         type: string
 
 concurrency:
-  group: test-cpu-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  group: test-cpu-python-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
   cancel-in-progress: true
 
 jobs:
-  test-cpu-no-tensor-engine:
-    name: Test CPU - No Tensor Engine
+  test-cpu-python:
+    name: Test CPU Python - No Tensor Engine
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       timeout: 60
@@ -38,3 +38,7 @@ jobs:
         # Tests requiring tensor engine / GPU need to be identified and flagged to skip.
         # We will just ensure monarch can be imported successfully.
         python -c "import monarch; print('Monarch imported successfully')"
+
+        # Run CPU Python tests (excluding GPU/tensor engine tests)
+        # TODO: Add appropriate test filters for CPU-only Python tests
+        echo "CPU Python tests would run here - currently placeholder"

.github/workflows/test-cuda.yml renamed to .github/workflows/test-gpu-python.yml

@@ -1,4 +1,4 @@
-name: Test CUDA
+name: Test GPU Python
 
 on:
   workflow_call:
@@ -9,12 +9,12 @@ on:
         type: string
 
 concurrency:
-  group: test-cuda-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  group: test-gpu-python-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
   cancel-in-progress: true
 
 jobs:
-  test-cuda:
-    name: Test CUDA (cuda12.6-py3.10)
+  test-gpu-python:
+    name: Test GPU Python (cuda12.6-py3.10)
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     strategy:
       fail-fast: true
@@ -39,10 +39,8 @@ jobs:
         # Setup test environment
         setup_test_environment
 
-        # Setup Tensor Engine dependencies
-        setup_tensor_engine
-
-        export CUDA_LIB_DIR=/usr/lib64
+        # Setup CUDA environment and library paths
+        setup_cuda_environment
 
         # Install the built wheel from artifact
         install_wheel_from_artifact
@@ -51,7 +49,7 @@ jobs:
         # pyre currently does not check these assertions
         pyright python/tests/test_python_actors.py
 
-        # Run CUDA tests
+        # Run GPU Python tests
         LC_ALL=C pytest python/tests/ -s -v -m "not oss_skip"
         # TODO(meriksen): temporarily disabled to unblock lands while debugging
         # mock CUDA issues on the OSS setup

.github/workflows/test-gpu-rust.yml

@@ -0,0 +1,60 @@
+name: Test GPU Rust
+
+on:
+  workflow_call:
+    inputs:
+      artifact-name:
+        description: 'Wheel artifact name from build workflow'
+        required: true
+        type: string
+
+concurrency:
+  group: test-gpu-rust-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test-gpu-rust:
+    name: Test GPU Rust (cuda12.6)
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
+    strategy:
+      fail-fast: true
+      matrix:
+        include:
+          - name: 4xlargegpu
+            runs-on: linux.g5.4xlarge.nvidia.gpu
+            torch-spec: '--pre torch --index-url https://download.pytorch.org/whl/nightly/cu126'
+            gpu-arch-type: "cuda"
+            gpu-arch-version: "12.6"
+    with:
+      timeout: 120
+      runner: ${{ matrix.runs-on }}
+      gpu-arch-type: ${{ matrix.gpu-arch-type }}
+      gpu-arch-version: ${{ matrix.gpu-arch-version }}
+      submodules: recursive
+      download-artifact: ${{ inputs.artifact-name }}
+      script: |
+        # Source common setup functions
+        source scripts/common-setup.sh
+
+        # Setup test environment
+        setup_test_environment
+
+        # Install Rust
+        setup_rust_toolchain
+
+        # Install System dependencies
+        install_system_dependencies
+
+        # Setup Tensor Engine dependencies
+        setup_tensor_engine
+
+        # Setup CUDA environment and library paths
+        setup_cuda_environment
+
+        # Setup PyTorch with C++ headers using common-setup utility
+        setup_pytorch_with_headers "${{ matrix.gpu-arch-version }}" "${{ matrix.torch-spec }}"
+
+        # Run GPU Rust tests
+        echo "Running OSS Rust tests..."
+        # TODO: fix broken tests, then update to `cargo test --no-fail-fast`
+        cargo test -p monarch_rdma

cuda-sys/build.rs

@@ -9,7 +9,7 @@
 use std::env;
 use std::path::PathBuf;
 
-use build_utils::*;
+use build_utils;
 
 #[cfg(target_os = "macos")]
 fn main() {}

hyperactor_macros/tests/export.rs

@@ -15,8 +15,7 @@ use hyperactor::Named;
 use hyperactor::PortRef;
 use hyperactor::Unbind;
 use serde::Deserialize;
-
-use crate::Serialize;
+use serde::Serialize;
 
 #[derive(Debug)]
 #[hyperactor::export(

hyperactor_telemetry/Cargo.toml

@@ -14,14 +14,14 @@ edition = "2024"
 anyhow = "1.0.98"
 chrono = { version = "0.4.41", features = ["clock", "serde", "std"], default-features = false }
 dashmap = { version = "5.5.3", features = ["rayon", "serde"] }
-fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main", optional = true }
+fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 hdrhistogram = "7.5"
 lazy_static = "1.5"
 opentelemetry = "0.29"
 opentelemetry_sdk = { version = "0.29.0", features = ["rt-tokio"] }
 rand = { version = "0.8", features = ["small_rng"] }
 rusqlite = { version = "0.36.0", features = ["backup", "blob", "bundled", "column_decltype", "functions", "limits", "modern_sqlite", "serde_json"] }
-scuba = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main", optional = true }
+scuba = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 serde_rusqlite = "0.39.3"
@@ -35,15 +35,9 @@ tracing-subscriber = { version = "0.3.20", features = ["chrono", "env-filter", "
 whoami = "1.5"
 
 [dev-dependencies]
-fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 quickcheck = "1.0"
 quickcheck_macros = "1.0"
-scuba = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 tracing-test = { version = "0.2.3", features = ["no-env-filter"] }
 
-[features]
-default = []
-fbcode_build = ["fbinit", "scuba"]
-
 [lints]
 rust = { unexpected_cfgs = { check-cfg = ["cfg(fbcode_build)"], level = "warn" } }

monarch_rdma/build.rs

@@ -51,10 +51,6 @@ fn main() {
     println!("cargo:rustc-link-lib=cuda");
     println!("cargo:rustc-link-lib=cudart");
 
-    // Tell cargo to look for shared libraries in standard RDMA directories
-    println!("cargo:rustc-link-search=/usr/lib");
-    println!("cargo:rustc-link-search=/usr/lib64");
-
     // Link against the ibverbs and mlx5 libraries (used by rdmaxcel-sys)
     println!("cargo:rustc-link-lib=ibverbs");
     println!("cargo:rustc-link-lib=mlx5");

monarch_tensor_worker/src/stream.rs

@@ -2200,11 +2200,12 @@ mod tests {
                 .unwrap()
                 .unwrap();
 
-            allclose(
+            let result = allclose(
                 &factory_float_tensor(data, "cpu".try_into().unwrap()),
                 &actual.borrow(),
             )
-            .unwrap()
+            .unwrap();
+            result
         }
 
         async fn validate_dependent_error(&mut self, reference: Ref, error: Arc<SeqError>) {

rdmaxcel-sys/build.rs

@@ -10,17 +10,11 @@ use std::env;
 use std::path::Path;
 use std::path::PathBuf;
 
-use build_utils::*;
-
 #[cfg(target_os = "macos")]
 fn main() {}
 
 #[cfg(not(target_os = "macos"))]
 fn main() {
-    // Tell cargo to look for shared libraries in the specified directory
-    println!("cargo:rustc-link-search=/usr/lib");
-    println!("cargo:rustc-link-search=/usr/lib64");
-
     // Link against the ibverbs library
     println!("cargo:rustc-link-lib=ibverbs");
 
@@ -141,8 +135,6 @@ fn main() {
     }
     if let Some(lib_dir) = &python_config.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);
     }