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.git
The following commit(s) were added to refs/heads/main by this push:
new 0ff0f85d1c [Runtime][Tests] Fix contrib wheel tests (#19714)
0ff0f85d1c is described below
commit 0ff0f85d1c8f3cce1a5968b2bf2dd590bd94b2dd
Author: Shushi Hong <[email protected]>
AuthorDate: Wed Jun 10 18:12:28 2026 -0400
[Runtime][Tests] Fix contrib wheel tests (#19714)
This pr fixes two contrib test failures that show up when running
source-tree tests against the `apache-tvm` wheel.
- Launch the pickle memoize helper script with `sys.executable` so
subprocesses use the same Python environment as pytest.
- Fix `tvm.contrib.random.random_fill` for packed sub-byte dtypes by
using the actual tensor storage size instead of the logical element
count.
## Root Cause
`tests/python/contrib/test_memoize.py` executed
`pickle_memoize_script.py` directly. In a wheel test environment, the
shebang can resolve to a different Python than the active wheel venv,
causing `ModuleNotFoundError: No module named 'tvm'`.
`random_fill` used the product of tensor shape as the number of values
to write. For packed dtypes such as `int4`, two logical elements share
one byte, so this wrote past the allocated storage and caused native
heap corruption / process abort.
---------
Co-authored-by: tqchen <[email protected]>
---
src/runtime/extra/contrib/random/mt_random_engine.cc | 16 +++-------------
tests/python/contrib/test_memoize.py | 8 ++++----
2 files changed, 7 insertions(+), 17 deletions(-)
diff --git a/src/runtime/extra/contrib/random/mt_random_engine.cc
b/src/runtime/extra/contrib/random/mt_random_engine.cc
index c01fe92673..d121bda59d 100644
--- a/src/runtime/extra/contrib/random/mt_random_engine.cc
+++ b/src/runtime/extra/contrib/random/mt_random_engine.cc
@@ -147,15 +147,7 @@ class RandomEngine {
// quantized dtype (uint8 / int8) data non-empty requirement
std::uniform_real_distribution<> dist(1.0, 10.0);
// Use float representation could make us work well on float / int type
too.
- if (dtype.bits == 1) {
- std::generate_n(static_cast<bool*>(data) + st, ed - st, [&]() { return
dist(rnd_engine_); });
- } else if (dtype.bits == 4) {
- // For uint4/int4 we pack two values into a single byte.
- // Thus, to ensure both values are non-zero, we use a distribution of 17
- 30.
- std::uniform_real_distribution<> packed_dist(17.0, 30.0);
- std::generate_n(reinterpret_cast<uint8_t*>(data) + st, ed - st,
- [&]() { return packed_dist(rnd_engine_); });
- } else if (dtype.bits == 8) {
+ if (dtype.bits == 8) {
std::generate_n(static_cast<uint8_t*>(data) + st, ed - st,
[&]() { return dist(rnd_engine_); });
} else if (dtype.bits == 16) {
@@ -180,8 +172,7 @@ class RandomEngine {
size *= tensor->shape[i];
}
DLDataType dtype = tensor->dtype;
- if (dtype.bits == 1 || dtype.bits == 4 || dtype.bits == 8 || dtype.bits ==
16 ||
- dtype.bits == 32 || dtype.bits == 64) {
+ if (dtype.bits == 8 || dtype.bits == 16 || dtype.bits == 32 || dtype.bits
== 64) {
FillDataImpl(tensor->data, 0, size, dtype);
} else {
TVM_FFI_THROW(InternalError)
@@ -218,8 +209,7 @@ class RandomEngine {
for (int i = 0; i < tensor->ndim; ++i) {
size *= tensor->shape[i];
}
- if (dtype.bits == 1 || dtype.bits == 4 || dtype.bits == 8 || dtype.bits ==
16 ||
- dtype.bits == 32 || dtype.bits == 64) {
+ if (dtype.bits == 8 || dtype.bits == 16 || dtype.bits == 32 || dtype.bits
== 64) {
int res = TVMBackendParallelLaunch(ParallelTask::RunTask, &task, 0);
TVM_FFI_ICHECK_EQ(res, 0) << "RandomFillForMeasure:
TVMBackendParallelLaunch failed";
} else {
diff --git a/tests/python/contrib/test_memoize.py
b/tests/python/contrib/test_memoize.py
index c3ebd73d41..e755b17521 100644
--- a/tests/python/contrib/test_memoize.py
+++ b/tests/python/contrib/test_memoize.py
@@ -31,7 +31,7 @@ TEST_SCRIPT_FILE =
pathlib.Path(__file__).with_name("pickle_memoize_script.py").
def test_cache_dir_not_in_current_working_dir():
with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
temp_dir = pathlib.Path(temp_dir)
- subprocess.check_call([TEST_SCRIPT_FILE, "1", "1"], cwd=temp_dir)
+ subprocess.check_call([sys.executable, str(TEST_SCRIPT_FILE), "1",
"1"], cwd=temp_dir)
new_files = list(temp_dir.iterdir())
assert not new_files, (
@@ -62,7 +62,7 @@ def test_cache_dir_defaults_to_home_config_cache():
with tempfile.TemporaryDirectory(prefix="tvm_") as temp_dir:
temp_dir = pathlib.Path(temp_dir)
- subprocess.check_call([TEST_SCRIPT_FILE, "1", "0"], cwd=temp_dir)
+ subprocess.check_call([sys.executable, str(TEST_SCRIPT_FILE), "1",
"0"], cwd=temp_dir)
new_files = list(temp_dir.iterdir())
assert not new_files, (
@@ -84,7 +84,7 @@ def test_cache_dir_respects_xdg_cache_home():
temp_working_dir = pathlib.Path(temp_working_dir)
subprocess.check_call(
- [TEST_SCRIPT_FILE, "1", "0"],
+ [sys.executable, str(TEST_SCRIPT_FILE), "1", "0"],
cwd=temp_working_dir,
env={
**os.environ,
@@ -112,7 +112,7 @@ def test_cache_dir_only_created_when_used():
temp_working_dir = pathlib.Path(temp_working_dir)
subprocess.check_call(
- [TEST_SCRIPT_FILE, "0", "1"],
+ [sys.executable, str(TEST_SCRIPT_FILE), "0", "1"],
cwd=temp_working_dir,
env={
**os.environ,
