From def1f62b8a368f67a7e8338f4a9d91157343c239 Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Tue, 23 Sep 2025 07:12:56 -0700 Subject: [PATCH 1/7] add back tests --- .github/workflows/gpu-tests.yml | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 8d9bfed..20abe8e 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -33,19 +33,20 @@ jobs: python -m pip install --upgrade pip python -m pip install --upgrade uv python -m uv pip install -U pytest "jax[cuda12]" - python -m uv pip install nvidia-cusolver-cu12==11.7.3.90 - python -m uv pip install nvidia-cublas-cu12 - python -m uv pip install jax-triton triton==3.3.1 + # python -m uv pip install nvidia-cusolver-cu12==11.7.3.90 + # python -m uv pip install nvidia-cublas-cu12 + # python -m uv pip install jax-triton triton==3.3.1 # python -m uv pip uninstall cuequivariance cuequivariance_jax cuequivariance_torch python -m uv pip install cuequivariance-ops-cu12 cuequivariance-ops-jax-cu12 + python -c "import cuequivariance_ops; print('cueop', cuequivariance_ops.__version__)" + python -c "import cuequivariance_ops_jax; print('cueopx', cuequivariance_ops_jax.__version__)" + python -m uv pip install -e ./cuequivariance python -m uv pip install -e ./cuequivariance_jax - - # python -c "import cuequivariance; print('cue', cuequivariance.__version__)" - # python -c "import cuequivariance_jax; print('cuex', cuequivariance_jax.__version__)" + python -c "import cuequivariance; print('cue', cuequivariance.__version__)" + python -c "import cuequivariance_jax; print('cuex', cuequivariance_jax.__version__)" - name: Test with pytest run: | - # XLA_PYTHON_CLIENT_PREALLOCATE=false pytest --doctest-modules -x -m "not slow" cuequivariance_jax - echo "skipping tests" + XLA_PYTHON_CLIENT_PREALLOCATE=false pytest --doctest-modules -x -m "not slow" cuequivariance_jax From a98f867fbf9d69607053ae78829ac8d8f96d9c52 Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Tue, 23 Sep 2025 07:27:44 -0700 Subject: [PATCH 2/7] add debug --- .github/workflows/gpu-tests.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 20abe8e..775f6e8 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -38,6 +38,14 @@ jobs: # python -m uv pip install jax-triton triton==3.3.1 # python -m uv pip uninstall cuequivariance cuequivariance_jax cuequivariance_torch python -m uv pip install cuequivariance-ops-cu12 cuequivariance-ops-jax-cu12 + + # add some debug print here, one liners please + echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH" + find $(python -c "import site; print(' '.join(site.getsitepackages()))") -name "*libcublas*" 2>/dev/null | head -5 + python -c "import os; print('CUDA_PATH:', os.environ.get('CUDA_PATH', 'Not set'))" + ldconfig -p | grep cublas || echo "No cublas in ldconfig" + python -c "import cuequivariance_ops; print('ops lib path:', cuequivariance_ops.__file__)" + python -c "import cuequivariance_ops; print('cueop', cuequivariance_ops.__version__)" python -c "import cuequivariance_ops_jax; print('cueopx', cuequivariance_ops_jax.__version__)" From e0c7032870ae68b4aa99c23fe34eb5c7dfaf2991 Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Tue, 23 Sep 2025 07:45:23 -0700 Subject: [PATCH 3/7] try --- .github/workflows/gpu-tests.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 775f6e8..f2767f5 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -45,6 +45,10 @@ jobs: python -c "import os; print('CUDA_PATH:', os.environ.get('CUDA_PATH', 'Not set'))" ldconfig -p | grep cublas || echo "No cublas in ldconfig" python -c "import cuequivariance_ops; print('ops lib path:', cuequivariance_ops.__file__)" + + # Add NVIDIA CUDA libraries to LD_LIBRARY_PATH + export LD_LIBRARY_PATH="$(python -c "import site; import os; print(':'.join([os.path.join(p, 'nvidia/cublas/lib') for p in site.getsitepackages() if os.path.exists(os.path.join(p, 'nvidia/cublas/lib'))]))"):$LD_LIBRARY_PATH + echo "Updated LD_LIBRARY_PATH: $LD_LIBRARY_PATH" python -c "import cuequivariance_ops; print('cueop', cuequivariance_ops.__version__)" python -c "import cuequivariance_ops_jax; print('cueopx', cuequivariance_ops_jax.__version__)" @@ -57,4 +61,6 @@ jobs: - name: Test with pytest run: | + # Set up CUDA library path for tests + export LD_LIBRARY_PATH="$(python -c "import site; import os; print(':'.join([os.path.join(p, 'nvidia/cublas/lib') for p in site.getsitepackages() if os.path.exists(os.path.join(p, 'nvidia/cublas/lib'))]))"):$LD_LIBRARY_PATH XLA_PYTHON_CLIENT_PREALLOCATE=false pytest --doctest-modules -x -m "not slow" cuequivariance_jax From 73b2a088b0706b231581b69b8bf6a0b125c76fbf Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Tue, 23 Sep 2025 07:55:15 -0700 Subject: [PATCH 4/7] try --- .github/workflows/gpu-tests.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index f2767f5..1e9f658 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -47,7 +47,7 @@ jobs: python -c "import cuequivariance_ops; print('ops lib path:', cuequivariance_ops.__file__)" # Add NVIDIA CUDA libraries to LD_LIBRARY_PATH - export LD_LIBRARY_PATH="$(python -c "import site; import os; print(':'.join([os.path.join(p, 'nvidia/cublas/lib') for p in site.getsitepackages() if os.path.exists(os.path.join(p, 'nvidia/cublas/lib'))]))"):$LD_LIBRARY_PATH + export LD_LIBRARY_PATH="$(python -c 'import site; import os; print(":".join([os.path.join(p, "nvidia/cublas/lib") for p in site.getsitepackages() if os.path.exists(os.path.join(p, "nvidia/cublas/lib"))]))'):$LD_LIBRARY_PATH echo "Updated LD_LIBRARY_PATH: $LD_LIBRARY_PATH" python -c "import cuequivariance_ops; print('cueop', cuequivariance_ops.__version__)" @@ -62,5 +62,5 @@ jobs: - name: Test with pytest run: | # Set up CUDA library path for tests - export LD_LIBRARY_PATH="$(python -c "import site; import os; print(':'.join([os.path.join(p, 'nvidia/cublas/lib') for p in site.getsitepackages() if os.path.exists(os.path.join(p, 'nvidia/cublas/lib'))]))"):$LD_LIBRARY_PATH + export LD_LIBRARY_PATH="$(python -c 'import site; import os; print(":".join([os.path.join(p, "nvidia/cublas/lib") for p in site.getsitepackages() if os.path.exists(os.path.join(p, "nvidia/cublas/lib"))]))'):$LD_LIBRARY_PATH XLA_PYTHON_CLIENT_PREALLOCATE=false pytest --doctest-modules -x -m "not slow" cuequivariance_jax From 20d3300f461e7f6166a0222f9dcd1d6da75ffe2a Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Tue, 23 Sep 2025 08:07:53 -0700 Subject: [PATCH 5/7] test --- .github/workflows/gpu-tests.yml | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 1e9f658..69d9ef0 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -44,13 +44,20 @@ jobs: find $(python -c "import site; print(' '.join(site.getsitepackages()))") -name "*libcublas*" 2>/dev/null | head -5 python -c "import os; print('CUDA_PATH:', os.environ.get('CUDA_PATH', 'Not set'))" ldconfig -p | grep cublas || echo "No cublas in ldconfig" + echo "=== About to check ops lib path ===" python -c "import cuequivariance_ops; print('ops lib path:', cuequivariance_ops.__file__)" + echo "=== About to find CUDA lib dir ===" # Add NVIDIA CUDA libraries to LD_LIBRARY_PATH - export LD_LIBRARY_PATH="$(python -c 'import site; import os; print(":".join([os.path.join(p, "nvidia/cublas/lib") for p in site.getsitepackages() if os.path.exists(os.path.join(p, "nvidia/cublas/lib"))]))'):$LD_LIBRARY_PATH + CUDA_LIB_DIR=$(find $(python -c "import site; print(' '.join(site.getsitepackages()))") -path "*/nvidia/cublas/lib" -type d | head -1) + echo "=== About to export LD_LIBRARY_PATH ===" + export LD_LIBRARY_PATH="$CUDA_LIB_DIR:$LD_LIBRARY_PATH" + echo "Added CUDA lib dir: $CUDA_LIB_DIR" echo "Updated LD_LIBRARY_PATH: $LD_LIBRARY_PATH" + echo "=== About to check cuequivariance_ops version ===" python -c "import cuequivariance_ops; print('cueop', cuequivariance_ops.__version__)" + echo "=== About to check cuequivariance_ops_jax version ===" python -c "import cuequivariance_ops_jax; print('cueopx', cuequivariance_ops_jax.__version__)" python -m uv pip install -e ./cuequivariance @@ -62,5 +69,6 @@ jobs: - name: Test with pytest run: | # Set up CUDA library path for tests - export LD_LIBRARY_PATH="$(python -c 'import site; import os; print(":".join([os.path.join(p, "nvidia/cublas/lib") for p in site.getsitepackages() if os.path.exists(os.path.join(p, "nvidia/cublas/lib"))]))'):$LD_LIBRARY_PATH + CUDA_LIB_DIR=$(find $(python -c "import site; print(' '.join(site.getsitepackages()))") -path "*/nvidia/cublas/lib" -type d | head -1) + export LD_LIBRARY_PATH="$CUDA_LIB_DIR:$LD_LIBRARY_PATH" XLA_PYTHON_CLIENT_PREALLOCATE=false pytest --doctest-modules -x -m "not slow" cuequivariance_jax From 18201f6458f72dc5fd2ed57db8b87747c000160d Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Tue, 23 Sep 2025 08:15:39 -0700 Subject: [PATCH 6/7] try --- .github/workflows/gpu-tests.yml | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 69d9ef0..34829fb 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -47,12 +47,13 @@ jobs: echo "=== About to check ops lib path ===" python -c "import cuequivariance_ops; print('ops lib path:', cuequivariance_ops.__file__)" - echo "=== About to find CUDA lib dir ===" - # Add NVIDIA CUDA libraries to LD_LIBRARY_PATH - CUDA_LIB_DIR=$(find $(python -c "import site; print(' '.join(site.getsitepackages()))") -path "*/nvidia/cublas/lib" -type d | head -1) + echo "=== About to find CUDA lib dirs ===" + # Add ALL NVIDIA CUDA libraries to LD_LIBRARY_PATH + SITE_PACKAGES=$(python -c "import site; print(' '.join(site.getsitepackages()))") + CUDA_LIB_DIRS=$(find $SITE_PACKAGES -path "*/nvidia/*/lib" -type d 2>/dev/null | tr '\n' ':') echo "=== About to export LD_LIBRARY_PATH ===" - export LD_LIBRARY_PATH="$CUDA_LIB_DIR:$LD_LIBRARY_PATH" - echo "Added CUDA lib dir: $CUDA_LIB_DIR" + export LD_LIBRARY_PATH="$CUDA_LIB_DIRS$LD_LIBRARY_PATH" + echo "Added CUDA lib dirs: $CUDA_LIB_DIRS" echo "Updated LD_LIBRARY_PATH: $LD_LIBRARY_PATH" echo "=== About to check cuequivariance_ops version ===" @@ -69,6 +70,7 @@ jobs: - name: Test with pytest run: | # Set up CUDA library path for tests - CUDA_LIB_DIR=$(find $(python -c "import site; print(' '.join(site.getsitepackages()))") -path "*/nvidia/cublas/lib" -type d | head -1) - export LD_LIBRARY_PATH="$CUDA_LIB_DIR:$LD_LIBRARY_PATH" + SITE_PACKAGES=$(python -c "import site; print(' '.join(site.getsitepackages()))") + CUDA_LIB_DIRS=$(find $SITE_PACKAGES -path "*/nvidia/*/lib" -type d 2>/dev/null | tr '\n' ':') + export LD_LIBRARY_PATH="$CUDA_LIB_DIRS$LD_LIBRARY_PATH" XLA_PYTHON_CLIENT_PREALLOCATE=false pytest --doctest-modules -x -m "not slow" cuequivariance_jax From da73ffba090c35af0218cf8d21008b98333c6f10 Mon Sep 17 00:00:00 2001 From: Mario Geiger Date: Wed, 24 Sep 2025 03:08:07 -0700 Subject: [PATCH 7/7] clean --- .github/workflows/gpu-tests.yml | 20 +------------------- 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index 34829fb..a640de5 100644 --- a/.github/workflows/gpu-tests.yml +++ b/.github/workflows/gpu-tests.yml @@ -33,32 +33,14 @@ jobs: python -m pip install --upgrade pip python -m pip install --upgrade uv python -m uv pip install -U pytest "jax[cuda12]" - # python -m uv pip install nvidia-cusolver-cu12==11.7.3.90 - # python -m uv pip install nvidia-cublas-cu12 - # python -m uv pip install jax-triton triton==3.3.1 - # python -m uv pip uninstall cuequivariance cuequivariance_jax cuequivariance_torch python -m uv pip install cuequivariance-ops-cu12 cuequivariance-ops-jax-cu12 - # add some debug print here, one liners please - echo "LD_LIBRARY_PATH: $LD_LIBRARY_PATH" - find $(python -c "import site; print(' '.join(site.getsitepackages()))") -name "*libcublas*" 2>/dev/null | head -5 - python -c "import os; print('CUDA_PATH:', os.environ.get('CUDA_PATH', 'Not set'))" - ldconfig -p | grep cublas || echo "No cublas in ldconfig" - echo "=== About to check ops lib path ===" - python -c "import cuequivariance_ops; print('ops lib path:', cuequivariance_ops.__file__)" - - echo "=== About to find CUDA lib dirs ===" - # Add ALL NVIDIA CUDA libraries to LD_LIBRARY_PATH + # Add NVIDIA CUDA libraries to LD_LIBRARY_PATH SITE_PACKAGES=$(python -c "import site; print(' '.join(site.getsitepackages()))") CUDA_LIB_DIRS=$(find $SITE_PACKAGES -path "*/nvidia/*/lib" -type d 2>/dev/null | tr '\n' ':') - echo "=== About to export LD_LIBRARY_PATH ===" export LD_LIBRARY_PATH="$CUDA_LIB_DIRS$LD_LIBRARY_PATH" - echo "Added CUDA lib dirs: $CUDA_LIB_DIRS" - echo "Updated LD_LIBRARY_PATH: $LD_LIBRARY_PATH" - echo "=== About to check cuequivariance_ops version ===" python -c "import cuequivariance_ops; print('cueop', cuequivariance_ops.__version__)" - echo "=== About to check cuequivariance_ops_jax version ===" python -c "import cuequivariance_ops_jax; print('cueopx', cuequivariance_ops_jax.__version__)" python -m uv pip install -e ./cuequivariance