Misc fixes for various build/test environments (#331)

* Allow the build to skip the TensorRT-based decoder if the user
explicitly disables it. (Useful for developer testing.)
* Fix a build warning from the UCCGSD test
* Also, skip the
`libs/solvers/python/tests/test_uccgsd.py::test_solvers_vqe_uccgsd_lih`
test in CPU-only environments (like some of our CI) because it was
taking ~30 minutes just for that one test.

Author: bmhowe23

docker/build_env/cudaqx.dev.Dockerfile

@@ -19,9 +19,21 @@ LABEL org.opencontainers.image.url="https://github.com/NVIDIA/cudaqx"
 # FIXME: Remove the cmake install once private repos are updated.
 RUN apt-get update && CUDA_DASH=$(echo $cuda_version | tr '.' '-') \
   && apt-get install -y gfortran libblas-dev jq cuda-nvtx-${CUDA_DASH} \
+  && apt-get install -y git-lfs \
   && python3 -m pip install "cmake<4" --user \
   && apt-get autoremove -y --purge && apt-get clean && rm -rf /var/lib/apt/lists/*
 
+# Install tensorrt-dev, unless this is CUDA 12 w/ arm64
+RUN CUDA_MAJOR_VERSION=$(echo $cuda_version | cut -d . -f1); \
+  ARCH="${ARCH:-$(dpkg --print-architecture)}"; \
+  if [ "$CUDA_MAJOR_VERSION" != "12" ] || [ "$ARCH" = "amd64" ]; then \
+    apt-get update \
+    && apt-cache search tensorrt | awk '{print "Package: "$1"\nPin: version *+cuda'${cuda_version}'\nPin-Priority: 1001\n"}' | tee /etc/apt/preferences.d/tensorrt-cuda${cuda_version}.pref > /dev/null \
+    && apt update \
+    && apt-get install -y tensorrt-dev \
+    && apt-get autoremove -y --purge && apt-get clean && rm -rf /var/lib/apt/lists/*; \
+  fi
+
 COPY .cudaq_version /cudaq_version
 
 ENV CUDAQ_INSTALL_PREFIX=/usr/local/cudaq

libs/qec/CMakeLists.txt

@@ -61,6 +61,9 @@ option(CUDAQX_QEC_INSTALL_PYTHON
        "Install python files alongside the library."
        ${CUDAQX_INSTALL_PYTHON})
 
+# Option to control TRT decoder build (default: ON)
+option(CUDAQ_QEC_BUILD_TRT_DECODER "Build the TensorRT decoder plugin" ON)
+
 # Check for CUDA Support (ref: cuda-quantum/CMakeLists.txt)
 # ==============================================================================
 include(CheckLanguage)

libs/qec/lib/CMakeLists.txt

@@ -25,7 +25,11 @@ add_library(${LIBRARY_NAME} SHARED
 )
 
 add_subdirectory(decoders/plugins/example)
-add_subdirectory(decoders/plugins/trt_decoder)
+
+if(CUDAQ_QEC_BUILD_TRT_DECODER)
+  add_subdirectory(decoders/plugins/trt_decoder)
+endif()
+
 add_subdirectory(codes)
 add_subdirectory(device) 
 

libs/qec/unittests/CMakeLists.txt

@@ -44,7 +44,7 @@ add_dependencies(CUDAQXQECUnitTests test_qec)
 gtest_discover_tests(test_qec)
 
 # TensorRT decoder test is only built for x86 architectures
-if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64)|(AMD64|amd64)|(^i.86$)")
+if(CUDAQ_QEC_BUILD_TRT_DECODER AND CMAKE_SYSTEM_PROCESSOR MATCHES "(x86_64)|(AMD64|amd64)|(^i.86$)")
   add_executable(test_trt_decoder ./decoders/trt_decoder/test_trt_decoder.cpp)
 
   # Find TensorRT for the test

libs/solvers/python/tests/test_uccgsd.py

@@ -6,6 +6,22 @@
 import cudaq
 import cudaq_solvers as solvers
 from scipy.optimize import minimize
+import subprocess
+
+
+def is_nvidia_gpu_available():
+    """Check if NVIDIA GPU is available using nvidia-smi command."""
+    try:
+        result = subprocess.run(["nvidia-smi"],
+                                stdout=subprocess.PIPE,
+                                stderr=subprocess.PIPE,
+                                text=True)
+        if result.returncode == 0 and "GPU" in result.stdout:
+            return True
+    except FileNotFoundError:
+        # The nvidia-smi command is not found, indicating no NVIDIA GPU drivers
+        return False
+    return False
 
 
 def test_solvers_adapt_uccgsd_h2():
@@ -138,6 +154,9 @@ def cost(theta):
     assert np.isclose(energy, -1.1371, atol=1e-4)
 
 
+# Since this test is so slow on the CPU, only run this test if a GPU was found.
+@pytest.mark.skipif(not is_nvidia_gpu_available(),
+                    reason="NVIDIA GPU not found")
 def test_solvers_vqe_uccgsd_lih():
 
     geometry = [('Li', (0.3925, 0., 0.)), ('H', (-1.1774, 0., .0))]

libs/solvers/unittests/test_uccgsd_operator_pool.cpp

@@ -187,8 +187,9 @@ TEST(UCCGSDOperatorPoolTest, AllOperatorsNonEmpty) {
   for (size_t i = 0; i < ops.size(); ++i) {
     EXPECT_GT(ops[i].num_terms(), 0)
         << "Operator " << i << " is empty (has 0 terms)";
-    EXPECT_FALSE(ops[i].is_identity())
-        << "Operator " << i << " should not be identity";
+    for (const auto &term : ops[i])
+      EXPECT_FALSE(term.is_identity())
+          << "Term " << term.to_string() << " is identity";
   }
 }