This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 86c45a54 [TEST] Run the Python test suite in parallel with a shared
GPU lock (#654)
86c45a54 is described below
commit 86c45a54ab72aa1d81b01ae86822fae4b959bd86
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Jul 5 23:57:11 2026 +0800
[TEST] Run the Python test suite in parallel with a shared GPU lock (#654)
---
pyproject.toml | 6 +-
python/tvm_ffi/testing/__init__.py | 1 +
python/tvm_ffi/testing/_locking.py | 121 +++++++++++++++++++++++++++
tests/python/test_cubin_launcher.py | 95 ++++++++++++---------
tests/python/test_current_work_stream_gpu.py | 17 ++--
tests/python/test_dlpack_exchange_api.py | 23 +++--
tests/python/test_load_inline.py | 26 ++++--
tests/python/test_stream.py | 94 ++++++++++++---------
tests/python/test_tensor.py | 13 ++-
9 files changed, 287 insertions(+), 109 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index cb3c4592..8cae0f0f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,7 +49,7 @@ torch = [
"numpy",
"ml_dtypes",
]
-test = [{ include-group = "torch" }, "pytest"]
+test = [{ include-group = "torch" }, "pytest", "pytest-xdist"]
dev = [
{ include-group = "test" },
"pre-commit",
@@ -176,6 +176,10 @@ sdist.exclude = [
[tool.pytest.ini_options]
testpaths = ["tests"]
+# Run the suite in parallel across all available cores. GPU-touching tests
+# serialize on a shared machine-local lock via
tvm_ffi.testing.run_with_gpu_lock,
+# so parallel workers do not contend for the device. Pass "-n0" to disable.
+addopts = ["-n", "auto"]
[tool.ruff]
include = ["python/**/*.py", "tests/**/*.py"]
diff --git a/python/tvm_ffi/testing/__init__.py
b/python/tvm_ffi/testing/__init__.py
index 4061bd63..22bcbb3b 100644
--- a/python/tvm_ffi/testing/__init__.py
+++ b/python/tvm_ffi/testing/__init__.py
@@ -17,6 +17,7 @@
"""Testing utilities."""
from ._ffi_api import * # noqa: F403
+from ._locking import run_with_gpu_lock
from .testing import (
TestCompare,
TestCustomCompare,
diff --git a/python/tvm_ffi/testing/_locking.py
b/python/tvm_ffi/testing/_locking.py
new file mode 100644
index 00000000..19a3f7e5
--- /dev/null
+++ b/python/tvm_ffi/testing/_locking.py
@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Helpers for tests that use exclusive machine-local resources.
+
+Parallel test runners (for example ``pytest -n auto``) spread tests across
+several worker processes. Tests that reach for a shared, machine-local
+resource such as a single GPU must serialize with each other so that
+concurrent access does not exhaust device memory or corrupt device state.
+:func:`run_with_gpu_lock` provides that serialization through an advisory
+file lock shared by every cooperating worker on the machine.
+"""
+
+from __future__ import annotations
+
+import getpass
+import os
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, TypeVar
+
+from ..utils import FileLock
+
+_LOCK_DIR_ENV_VAR = "TVM_FFI_TEST_LOCK_DIR"
+_LOCK_DIR_PREFIX = "tvm-ffi-test-locks"
+_LOCK_FILENAME = "gpu.lock"
+_R = TypeVar("_R")
+
+# Resolved GPU lock path, cached on the first ``run_with_gpu_lock`` call.
+#
+# Init is lazy (not at import) because importing does not imply a GPU test will
+# run, and resolving the path creates a directory. Concurrent first-callers may
+# race to set this global, but the race is benign: every worker derives the
same
+# path.
+_GPU_LOCK_PATH: Path | None = None
+
+
+def _ensure_gpu_lock_path() -> Path:
+ """Return the path to the machine-local GPU lock file, creating its
directory.
+
+ Returns
+ -------
+ path
+ The full path to the ``gpu.lock`` file. The parent directory is created
+ if it does not yet exist.
+
+ Notes
+ -----
+ The lock directory defaults to a per-user directory under the system
+ temporary directory, ``<tempdir>/tvm-ffi-test-locks-<user>``. Scoping the
+ default to the current user avoids ownership and permission conflicts when
+ several users share one host. It can be redirected with the
+ ``TVM_FFI_TEST_LOCK_DIR`` environment variable when all cooperating
+ processes need an explicitly shared machine-local path.
+
+ """
+ lock_dir_override = os.environ.get(_LOCK_DIR_ENV_VAR)
+ if lock_dir_override:
+ lock_dir = Path(lock_dir_override).expanduser()
+ else:
+ # Tag the default directory with the current user, falling back to the
+ # numeric uid then ``unknown`` when a login name cannot be resolved.
+ try:
+ user_tag = getpass.getuser()
+ except Exception:
+ uid = getattr(os, "getuid", None)
+ user_tag = str(uid()) if uid is not None else "unknown"
+ lock_dir = Path(tempfile.gettempdir()) /
f"{_LOCK_DIR_PREFIX}-{user_tag}"
+
+ lock_dir.mkdir(parents=True, exist_ok=True)
+ return lock_dir / _LOCK_FILENAME
+
+
+def run_with_gpu_lock(func: Callable[..., _R], /, *args: Any, **kwargs: Any)
-> _R:
+ """Run a callable while holding the machine-local GPU lock.
+
+ The lock serializes GPU access across parallel test workers so that
+ concurrent device use does not break GPU-related tests. Pass a callable
+ that contains the complete live-device lifetime (device creation,
+ allocation, execution, synchronization, and result checks); keep work that
+ does not touch the device, such as source compilation, outside the callable
+ so it can still run in parallel.
+
+ Parameters
+ ----------
+ func
+ Callable containing the complete live local-GPU lifetime.
+ args
+ Positional arguments forwarded to ``func``.
+ kwargs
+ Keyword arguments forwarded to ``func``.
+
+ Returns
+ -------
+ result
+ The return value of ``func``.
+
+ """
+ # Resolve and cache the lock path on the first call (see ``_GPU_LOCK_PATH``
+ # above for why this is lazy and why the concurrent-init race is benign).
+ global _GPU_LOCK_PATH # noqa: PLW0603 -- intentional first-call
memoization cache
+ lock_path = _GPU_LOCK_PATH
+ if lock_path is None:
+ lock_path = _ensure_gpu_lock_path()
+ _GPU_LOCK_PATH = lock_path
+ with FileLock(str(lock_path)):
+ return func(*args, **kwargs)
diff --git a/tests/python/test_cubin_launcher.py
b/tests/python/test_cubin_launcher.py
index d6f4e35b..bc038d09 100644
--- a/tests/python/test_cubin_launcher.py
+++ b/tests/python/test_cubin_launcher.py
@@ -32,6 +32,7 @@ except ImportError:
torch = None # ty: ignore[invalid-assignment]
import tvm_ffi.cpp
+from tvm_ffi.testing import run_with_gpu_lock
# Check if CUDA is available
@@ -201,28 +202,32 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_mul_two,
cubin_test::LaunchMulTwo);
extra_ldflags=["-lcudart"],
)
- # Load CUBIN from bytes
- load_fn = mod["load_cubin_data"]
- load_fn(cubin_bytes)
+ def run_and_check() -> None:
+ assert torch is not None
+ # Load CUBIN from bytes
+ load_fn = mod["load_cubin_data"]
+ load_fn(cubin_bytes)
- # Test add_one kernel
- launch_add_one = mod["launch_add_one"]
- n = 256
- x = torch.arange(n, dtype=torch.float32, device="cuda")
- y = torch.empty(n, dtype=torch.float32, device="cuda")
+ # Test add_one kernel
+ launch_add_one = mod["launch_add_one"]
+ n = 256
+ x = torch.arange(n, dtype=torch.float32, device="cuda")
+ y = torch.empty(n, dtype=torch.float32, device="cuda")
- launch_add_one(x, y)
- expected = x + 1
- torch.testing.assert_close(y, expected)
+ launch_add_one(x, y)
+ expected = x + 1
+ torch.testing.assert_close(y, expected)
- # Test mul_two kernel
- launch_mul_two = mod["launch_mul_two"]
- x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5
- y = torch.empty(n, dtype=torch.float32, device="cuda")
+ # Test mul_two kernel
+ launch_mul_two = mod["launch_mul_two"]
+ x = torch.arange(n, dtype=torch.float32, device="cuda") * 0.5
+ y = torch.empty(n, dtype=torch.float32, device="cuda")
- launch_mul_two(x, y)
- expected = x * 2
- torch.testing.assert_close(y, expected)
+ launch_mul_two(x, y)
+ expected = x * 2
+ torch.testing.assert_close(y, expected)
+
+ run_with_gpu_lock(run_and_check)
@pytest.mark.skipif(sys.platform != "linux", reason="CUBIN launcher only
supported on Linux")
@@ -299,17 +304,21 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_add_one_ex,
cubin_test_launch_ex::LaunchAdd
extra_ldflags=["-lcudart"],
)
- load_fn = mod["load_cubin_data"]
- load_fn(cubin_bytes)
+ def run_and_check() -> None:
+ assert torch is not None
+ load_fn = mod["load_cubin_data"]
+ load_fn(cubin_bytes)
+
+ launch_add_one_ex = mod["launch_add_one_ex"]
+ n = 256
+ x = torch.arange(n, dtype=torch.float32, device="cuda")
+ y = torch.empty(n, dtype=torch.float32, device="cuda")
- launch_add_one_ex = mod["launch_add_one_ex"]
- n = 256
- x = torch.arange(n, dtype=torch.float32, device="cuda")
- y = torch.empty(n, dtype=torch.float32, device="cuda")
+ launch_add_one_ex(x, y)
+ expected = x + 1
+ torch.testing.assert_close(y, expected)
- launch_add_one_ex(x, y)
- expected = x + 1
- torch.testing.assert_close(y, expected)
+ run_with_gpu_lock(run_and_check)
@pytest.mark.skipif(sys.platform != "linux", reason="CUBIN launcher only
supported on Linux")
@@ -387,21 +396,25 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(launch_mul_two,
cubin_test_chain::LaunchMulTwo);
mod = tvm_ffi.cpp.load_inline("cubin_test_chain", cuda_sources=cpp_code)
- # Load CUBIN from bytes
- load_fn = mod["load_cubin_data"]
- load_fn(cubin_bytes)
+ def run_and_check() -> None:
+ assert torch is not None
+ # Load CUBIN from bytes
+ load_fn = mod["load_cubin_data"]
+ load_fn(cubin_bytes)
+
+ launch_add_one = mod["launch_add_one"]
+ launch_mul_two = mod["launch_mul_two"]
- launch_add_one = mod["launch_add_one"]
- launch_mul_two = mod["launch_mul_two"]
+ # Test chained execution: (x + 1) * 2
+ n = 128
+ x = torch.full((n,), 5.0, dtype=torch.float32, device="cuda")
+ temp = torch.empty(n, dtype=torch.float32, device="cuda")
+ y = torch.empty(n, dtype=torch.float32, device="cuda")
- # Test chained execution: (x + 1) * 2
- n = 128
- x = torch.full((n,), 5.0, dtype=torch.float32, device="cuda")
- temp = torch.empty(n, dtype=torch.float32, device="cuda")
- y = torch.empty(n, dtype=torch.float32, device="cuda")
+ launch_add_one(x, temp) # temp = x + 1 = 6
+ launch_mul_two(temp, y) # y = temp * 2 = 12
- launch_add_one(x, temp) # temp = x + 1 = 6
- launch_mul_two(temp, y) # y = temp * 2 = 12
+ expected = torch.full((n,), 12.0, dtype=torch.float32, device="cuda")
+ torch.testing.assert_close(y, expected)
- expected = torch.full((n,), 12.0, dtype=torch.float32, device="cuda")
- torch.testing.assert_close(y, expected)
+ run_with_gpu_lock(run_and_check)
diff --git a/tests/python/test_current_work_stream_gpu.py
b/tests/python/test_current_work_stream_gpu.py
index 910962e0..aed771a3 100644
--- a/tests/python/test_current_work_stream_gpu.py
+++ b/tests/python/test_current_work_stream_gpu.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import ctypes
import pytest
+from tvm_ffi.testing import run_with_gpu_lock
try:
import torch
@@ -88,9 +89,13 @@ def test_current_work_stream_matches_torch_stream() -> None:
extra_include_paths=include_paths,
)
- device_id = torch.cuda.current_device()
- is_hip = torch.version.hip is not None
- stream = torch.cuda.Stream(device=device_id)
- with torch.cuda.stream(stream):
- expected_stream = int(stream.cuda_stream)
- mod.assert_current_work_stream(api_ptr, is_hip, expected_stream)
+ def run_and_check() -> None:
+ assert torch is not None
+ device_id = torch.cuda.current_device()
+ is_hip = torch.version.hip is not None
+ stream = torch.cuda.Stream(device=device_id)
+ with torch.cuda.stream(stream):
+ expected_stream = int(stream.cuda_stream)
+ mod.assert_current_work_stream(api_ptr, is_hip, expected_stream)
+
+ run_with_gpu_lock(run_and_check)
diff --git a/tests/python/test_dlpack_exchange_api.py
b/tests/python/test_dlpack_exchange_api.py
index 0938a251..f37e3b98 100644
--- a/tests/python/test_dlpack_exchange_api.py
+++ b/tests/python/test_dlpack_exchange_api.py
@@ -31,6 +31,7 @@ try:
import tvm_ffi
from torch.utils import cpp_extension
from tvm_ffi import libinfo
+ from tvm_ffi.testing import run_with_gpu_lock
except ImportError:
torch = None # ty: ignore[invalid-assignment]
@@ -223,17 +224,21 @@ def test_dlpack_exchange_api_gpu_tensor_metadata() ->
None:
assert torch is not None
echo = tvm_ffi.get_global_func("testing.echo")
- for shape in [(512,), (512, 512), (2, 3, 4)]:
- source = torch.empty(shape, device="cuda", dtype=torch.float16)
+ def run_and_check() -> None:
+ assert torch is not None
+ for shape in [(512,), (512, 512), (2, 3, 4)]:
+ source = torch.empty(shape, device="cuda", dtype=torch.float16)
- tvm_tensor = tvm_ffi.from_dlpack(source)
- assert tvm_tensor.shape == shape
- assert tvm_tensor.dtype == tvm_ffi.dtype("float16")
+ tvm_tensor = tvm_ffi.from_dlpack(source)
+ assert tvm_tensor.shape == shape
+ assert tvm_tensor.dtype == tvm_ffi.dtype("float16")
- echoed = echo(source)
- assert tuple(echoed.shape) == shape
- assert echoed.dtype == source.dtype
- assert echoed.device == source.device
+ echoed = echo(source)
+ assert tuple(echoed.shape) == shape
+ assert echoed.dtype == source.dtype
+ assert echoed.device == source.device
+
+ run_with_gpu_lock(run_and_check)
@pytest.mark.skipif(not _has_dlpack_api, reason="PyTorch DLPack Exchange API
not available")
diff --git a/tests/python/test_load_inline.py b/tests/python/test_load_inline.py
index 672a264c..3cc62da5 100644
--- a/tests/python/test_load_inline.py
+++ b/tests/python/test_load_inline.py
@@ -27,6 +27,7 @@ except ImportError:
import tvm_ffi.cpp
from tvm_ffi.module import Module
+from tvm_ffi.testing import run_with_gpu_lock
def test_load_inline_cpp() -> None:
@@ -191,7 +192,8 @@ def test_load_inline_cuda() -> None:
functions=["add_one_cuda"],
)
- if torch is not None:
+ def run_and_check() -> None:
+ assert torch is not None
# test with raw stream
x_cuda = torch.asarray([1, 2, 3, 4, 5], dtype=torch.float32,
device="cuda")
y_cuda = torch.empty_like(x_cuda)
@@ -206,6 +208,8 @@ def test_load_inline_cuda() -> None:
stream.synchronize()
torch.testing.assert_close(x_cuda + 1, y_cuda)
+ run_with_gpu_lock(run_and_check)
+
@pytest.mark.skipif(torch is None, reason="Requires torch")
def test_load_inline_with_env_tensor_allocator() -> None:
@@ -242,7 +246,7 @@ def test_load_inline_with_env_tensor_allocator() -> None:
)
assert torch is not None
- def run_check() -> None:
+ def run_and_check() -> None:
"""Must run in a separate function to ensure deletion happens before
mod unloads.
When a module returns an object, the object deleter address is part of
the
@@ -257,7 +261,7 @@ def test_load_inline_with_env_tensor_allocator() -> None:
assert y_cpu.dtype == torch.float32
torch.testing.assert_close(x_cpu + 1, y_cpu)
- run_check()
+ run_and_check()
@pytest.mark.skipif(
@@ -321,10 +325,14 @@ def test_load_inline_both() -> None:
mod.add_one_cpu(x, y)
numpy.testing.assert_equal(x + 1, y)
- x_cuda = torch.asarray([1, 2, 3, 4, 5], dtype=torch.float32, device="cuda")
- y_cuda = torch.empty_like(x_cuda)
- mod.add_one_cuda(x_cuda, y_cuda)
- torch.testing.assert_close(x_cuda + 1, y_cuda)
+ def run_and_check() -> None:
+ assert torch is not None
+ x_cuda = torch.asarray([1, 2, 3, 4, 5], dtype=torch.float32,
device="cuda")
+ y_cuda = torch.empty_like(x_cuda)
+ mod.add_one_cuda(x_cuda, y_cuda)
+ torch.testing.assert_close(x_cuda + 1, y_cuda)
+
+ run_with_gpu_lock(run_and_check)
@pytest.mark.skipif(
@@ -349,7 +357,7 @@ def test_cuda_memory_alloc_noleak() -> None:
functions=["return_tensor"],
)
- def run_check() -> None:
+ def run_and_check() -> None:
"""Must run in a separate function to ensure deletion happens before
mod unloads."""
assert torch is not None
x = torch.arange(1024 * 1024, dtype=torch.float32, device="cuda")
@@ -361,4 +369,4 @@ def test_cuda_memory_alloc_noleak() -> None:
# memory should not grow as we loop over
assert diff <= 1024**2 * 8
- run_check()
+ run_with_gpu_lock(run_and_check)
diff --git a/tests/python/test_stream.py b/tests/python/test_stream.py
index afc9139b..ffdd122e 100644
--- a/tests/python/test_stream.py
+++ b/tests/python/test_stream.py
@@ -22,6 +22,7 @@ import ctypes
import pytest
import tvm_ffi
import tvm_ffi.cpp
+from tvm_ffi.testing import run_with_gpu_lock
try:
import torch
@@ -83,21 +84,26 @@ def test_raw_stream() -> None:
def test_torch_stream() -> None:
assert torch is not None
mod = gen_check_stream_mod()
- device_id = torch.cuda.current_device()
- device = tvm_ffi.device("cuda", device_id)
- device_type = device.dlpack_device_type()
- stream_1 = torch.cuda.Stream(device_id)
- stream_2 = torch.cuda.Stream(device_id)
- with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_1)):
- assert torch.cuda.current_stream() == stream_1
- mod.check_stream(device_type, device_id, stream_1.cuda_stream)
- with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_2)):
- assert torch.cuda.current_stream() == stream_2
- mod.check_stream(device_type, device_id, stream_2.cuda_stream)
+ def run_and_check() -> None:
+ assert torch is not None
+ device_id = torch.cuda.current_device()
+ device = tvm_ffi.device("cuda", device_id)
+ device_type = device.dlpack_device_type()
+ stream_1 = torch.cuda.Stream(device_id)
+ stream_2 = torch.cuda.Stream(device_id)
+ with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_1)):
+ assert torch.cuda.current_stream() == stream_1
+ mod.check_stream(device_type, device_id, stream_1.cuda_stream)
+
+ with tvm_ffi.use_torch_stream(torch.cuda.stream(stream_2)):
+ assert torch.cuda.current_stream() == stream_2
+ mod.check_stream(device_type, device_id, stream_2.cuda_stream)
- assert torch.cuda.current_stream() == stream_1
- mod.check_stream(device_type, device_id, stream_1.cuda_stream)
+ assert torch.cuda.current_stream() == stream_1
+ mod.check_stream(device_type, device_id, stream_1.cuda_stream)
+
+ run_with_gpu_lock(run_and_check)
@pytest.mark.skipif(
@@ -106,24 +112,29 @@ def test_torch_stream() -> None:
def test_torch_current_stream() -> None:
assert torch is not None
mod = gen_check_stream_mod()
- device_id = torch.cuda.current_device()
- device = tvm_ffi.device("cuda", device_id)
- device_type = device.dlpack_device_type()
- stream_1 = torch.cuda.Stream(device_id)
- stream_2 = torch.cuda.Stream(device_id)
- with torch.cuda.stream(stream_1):
- assert torch.cuda.current_stream() == stream_1
- with tvm_ffi.use_torch_stream():
- mod.check_stream(device_type, device_id, stream_1.cuda_stream)
- with torch.cuda.stream(stream_2):
- assert torch.cuda.current_stream() == stream_2
+ def run_and_check() -> None:
+ assert torch is not None
+ device_id = torch.cuda.current_device()
+ device = tvm_ffi.device("cuda", device_id)
+ device_type = device.dlpack_device_type()
+ stream_1 = torch.cuda.Stream(device_id)
+ stream_2 = torch.cuda.Stream(device_id)
+ with torch.cuda.stream(stream_1):
+ assert torch.cuda.current_stream() == stream_1
with tvm_ffi.use_torch_stream():
- mod.check_stream(device_type, device_id, stream_2.cuda_stream)
+ mod.check_stream(device_type, device_id, stream_1.cuda_stream)
- assert torch.cuda.current_stream() == stream_1
- with tvm_ffi.use_torch_stream():
- mod.check_stream(device_type, device_id, stream_1.cuda_stream)
+ with torch.cuda.stream(stream_2):
+ assert torch.cuda.current_stream() == stream_2
+ with tvm_ffi.use_torch_stream():
+ mod.check_stream(device_type, device_id,
stream_2.cuda_stream)
+
+ assert torch.cuda.current_stream() == stream_1
+ with tvm_ffi.use_torch_stream():
+ mod.check_stream(device_type, device_id, stream_1.cuda_stream)
+
+ run_with_gpu_lock(run_and_check)
@pytest.mark.skipif(
@@ -132,14 +143,19 @@ def test_torch_current_stream() -> None:
def test_torch_graph() -> None:
assert torch is not None
mod = gen_check_stream_mod()
- device_id = torch.cuda.current_device()
- device = tvm_ffi.device("cuda", device_id)
- device_type = device.dlpack_device_type()
- graph = torch.cuda.CUDAGraph()
- stream = torch.cuda.Stream(device_id)
- x = torch.zeros(1, device="cuda")
- with tvm_ffi.use_torch_stream(torch.cuda.graph(graph, stream=stream)):
- assert torch.cuda.current_stream() == stream
- mod.check_stream(device_type, device_id, stream.cuda_stream)
- # avoid cuda graph no capture warning
- x = x + 1
+
+ def run_and_check() -> None:
+ assert torch is not None
+ device_id = torch.cuda.current_device()
+ device = tvm_ffi.device("cuda", device_id)
+ device_type = device.dlpack_device_type()
+ graph = torch.cuda.CUDAGraph()
+ stream = torch.cuda.Stream(device_id)
+ x = torch.zeros(1, device="cuda")
+ with tvm_ffi.use_torch_stream(torch.cuda.graph(graph, stream=stream)):
+ assert torch.cuda.current_stream() == stream
+ mod.check_stream(device_type, device_id, stream.cuda_stream)
+ # avoid cuda graph no capture warning
+ x = x + 1
+
+ run_with_gpu_lock(run_and_check)
diff --git a/tests/python/test_tensor.py b/tests/python/test_tensor.py
index b89b236a..e21de826 100644
--- a/tests/python/test_tensor.py
+++ b/tests/python/test_tensor.py
@@ -30,6 +30,7 @@ except ImportError:
import numpy as np
import tvm_ffi
+from tvm_ffi.testing import run_with_gpu_lock
def test_tensor_attributes() -> None:
@@ -194,10 +195,14 @@ def test_tensor_from_pytorch_rocm() -> None:
def _check_device(x: tvm_ffi.Tensor) -> str:
return x.device.type
- # PyTorch uses device name "cuda" to represent ROCm device
- x = torch.randn(128, device="cuda")
- device_type = tvm_ffi.get_global_func("testing.check_device")(x)
- assert device_type == "rocm"
+ def run_and_check() -> None:
+ assert torch is not None
+ # PyTorch uses device name "cuda" to represent ROCm device
+ x = torch.randn(128, device="cuda")
+ device_type = tvm_ffi.get_global_func("testing.check_device")(x)
+ assert device_type == "rocm"
+
+ run_with_gpu_lock(run_and_check)
def test_optional_tensor_view() -> None:
