Fix tests

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

.ci/scripts/test_llama.sh

@@ -58,6 +58,7 @@ cmake_install_executorch_libraries() {
     fi
     retry cmake -DBUCK2="$BUCK" \
         -DCMAKE_INSTALL_PREFIX=cmake-out \
+        -DCMAKE_PREFIX_PATH=$($PYTHON_EXECUTABLE -c "from distutils.sysconfig import get_python_lib; get_python_lib")
         -DCMAKE_BUILD_TYPE=Release \
         -DEXECUTORCH_BUILD_CUSTOM=ON \
         -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \

examples/models/llama2/custom_ops/op_sdpa_with_kv_cache_test.cpp

@@ -8,7 +8,7 @@
 
 #include <limits>
 
-#include <executorch/examples/models/llama2/custom_ops/FunctionHeaderWrapper.h> // Declares the operator
+#include <executorch/examples/models/llama2/custom_ops/op_sdpa.h> // Declares the operator
 #include <executorch/kernels/test/TestUtil.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>

examples/models/llama2/custom_ops/sdpa_with_kv_cache.py

@@ -9,6 +9,7 @@
 # This is only needed for OSS.
 
 import os
+from pathlib import Path
 
 import torch
 
@@ -17,8 +18,12 @@
     assert op is not None
 except:
     # assuming we only hit this in OSS, find the default install path
-    prefix = os.environ.get("CMAKE_INSTALL_PREFIX", "../../../../cmake-out")
-    lib_path = os.path.join(prefix, "lib/libcustom_ops_aot_lib.so")
+    if "CMAKE_INSTALL_PREFIX" in os.environ:
+        prefix = Path(os.environ["CMAKE_INSTALL_PREFIX"])
+    else:
+        prefix = Path(__file__).parent.parent.parent.parent.parent.resolve() / "cmake-out"
+    lib_path = prefix / "lib" / "libcustom_ops_aot_lib.so"
+    assert lib_path.exists(), f"{lib_path} does not exist, please run cmake --install"
     torch.ops.load_library(lib_path)
     op = torch.ops.llama.sdpa_with_kv_cache.default
     assert op is not None