fix tests

Author: yaoyaoding

python/tvm_ffi/_optional_torch_c_dlpack.py

@@ -60,17 +60,26 @@ def load_torch_c_dlpack_extension() -> Any:
 
         # check whether a JIT shared library is built in cache
         cache_dir = Path(os.environ.get("TVM_FFI_CACHE_DIR", "~/.cache/tvm-ffi")).expanduser()
-        addon_build_dir = cache_dir / "torch_c_dlpack_addon"
-        lib_path = addon_build_dir / (
-            "libtorch_c_dlpack_addon" + (".dll" if sys.platform == "win32" else ".so")
-        )
+        addon_output_dir = cache_dir
+        major, minor = torch.__version__.split(".")[:2]
+        device = "cpu" if not torch.cuda.is_available() else "cuda"
+        suffix = ".dll" if sys.platform.startswith("win") else ".so"
+        libname = f"libtorch_c_dlpack_addon_torch{major}{minor}-{device}{suffix}"
+        lib_path = addon_output_dir / libname
         if not lib_path.exists():
             build_script_path = (
                 Path(__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py"
             )
-            args = [sys.executable, str(build_script_path), "--build_dir", str(addon_build_dir)]
+            args = [
+                sys.executable,
+                str(build_script_path),
+                "--output-dir",
+                str(cache_dir),
+                "--libname",
+                libname,
+            ]
             if torch.cuda.is_available():
-                args.append("--build_with_cuda")
+                args.append("--build-with-cuda")
             subprocess.run(
                 args,
                 check=True,

tests/python/test_optional_torch_c_dlpack.py

@@ -36,15 +36,19 @@
 @pytest.mark.skipif(torch is None, reason="torch is not installed")
 def test_build_torch_c_dlpack_extension() -> None:
     build_script = Path(tvm_ffi.__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py"
-    subprocess.run(
-        [sys.executable, str(build_script), "--build_dir", "./build_test_dir"], check=True
-    )
+    args = [
+        sys.executable,
+        str(build_script),
+        "--output-dir",
+        "./output-dir",
+        "--libname",
+        "libtorch_c_dlpack_addon_test.so",
+    ]
+    if torch.cuda.is_available():
+        args.append("--build-with-cuda")
+    subprocess.run(args, check=True)
 
-    lib_path = str(
-        Path(
-            "./build_test_dir/libtorch_c_dlpack_addon.{}".format("dll" if IS_WINDOWS else "so")
-        ).resolve()
-    )
+    lib_path = str(Path("./output-dir/libtorch_c_dlpack_addon_test.so").resolve())
     assert Path(lib_path).exists()
 
     lib = ctypes.CDLL(lib_path)
@@ -58,20 +62,19 @@ def test_build_torch_c_dlpack_extension() -> None:
 def test_parallel_build() -> None:
     build_script = Path(tvm_ffi.__file__).parent / "utils" / "_build_optional_torch_c_dlpack.py"
     num_processes = 4
-    build_dir = "./build_test_dir_parallel"
+    output_dir = "./output-dir-parallel"
+    libname = "libtorch_c_dlpack_addon_test.so"
     processes = []
     for i in range(num_processes):
-        p = subprocess.Popen([sys.executable, str(build_script), "--build_dir", build_dir])
-        processes.append((p, build_dir))
+        p = subprocess.Popen(
+            [sys.executable, str(build_script), "--output-dir", output_dir, "--libname", libname]
+        )
+        processes.append((p, output_dir))
 
-    for p, build_dir in processes:
+    for p, output_dir in processes:
         p.wait()
         assert p.returncode == 0
-    lib_path = str(
-        Path(
-            "{}/libtorch_c_dlpack_addon.{}".format(build_dir, "dll" if IS_WINDOWS else "so")
-        ).resolve()
-    )
+    lib_path = str(Path(f"{output_dir}/{libname}").resolve())
     assert Path(lib_path).exists()