Shushi Hong
a979b2f98c
### Motivation `tvm.testing` imports `pytest` at module load (`tvm/testing/utils.py`). `tvm.rpc.server` imports `tvm.rpc.testing` (to register the `rpc.test.*` helpers), and `tvm.rpc.testing` imported `tvm.testing` at the top level, so a plain `import tvm` / `import tvm.relax` pulls `pytest` in through: ``` tvm.relax -> tvm.runtime.vm -> tvm.rpc -> rpc.server -> rpc.testing -> tvm.testing -> pytest ``` As a result `pytest` is effectively a runtime dependency: a user who installs TVM without `pytest` hits `ModuleNotFoundError: No module named 'pytest'` on import. This is easy to miss because test environments install `pytest`. ### Change `tvm.rpc.testing` only uses `tvm.testing.object_use_count` in a single test helper, so import it lazily at the call site instead of at module top level. This keeps the `rpc.test.*` registration and the helper behavior intact while removing `tvm.testing` (and `pytest`) from the `import tvm` path, so `pytest` can remain a test-only dependency. No functional change; `rpc.testing` is still imported by `rpc.server` and still registers the same global functions.
77 lines
2.3 KiB
Python
77 lines
2.3 KiB
Python
# 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.
|
|
|
|
# pylint: disable=invalid-name,unnecessary-comprehension
|
|
"""Testing functions for the RPC server."""
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
|
|
|
|
# RPC test functions to be registered for unit-tests purposes
|
|
@tvm.register_global_func("rpc.test.addone")
|
|
def _addone(x):
|
|
return x + 1
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.strcat")
|
|
def _strcat(name, x):
|
|
return f"{name}:{x}"
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.except")
|
|
def _remotethrow(name):
|
|
raise ValueError(f"{name}")
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.runtime_str_concat")
|
|
def _strcat(x, y):
|
|
return x + y
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.remote_tensor_func")
|
|
def _remote_tensor_func(y):
|
|
x = np.ones((3, 4))
|
|
np.testing.assert_equal(y.numpy(), x)
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.add_to_lhs")
|
|
def _add_to_lhs(x):
|
|
return lambda y: x + y
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.remote_return_nd")
|
|
def _my_module(name):
|
|
# Use closure to check the ref counter correctness
|
|
nd = tvm.runtime.tensor(np.zeros(10).astype("float32"))
|
|
|
|
if name == "get_arr":
|
|
return lambda: nd
|
|
if name == "ref_count":
|
|
# Imported lazily: rpc.server imports this module to register the
|
|
# rpc.test.* helpers, so a top-level ``import tvm.testing`` would drag
|
|
# tvm.testing (which imports pytest) into the plain ``import tvm`` path.
|
|
from tvm.testing import object_use_count
|
|
|
|
return lambda: object_use_count(nd)
|
|
if name == "get_elem":
|
|
return lambda idx: nd.numpy()[idx]
|
|
if name == "get_arr_elem":
|
|
return lambda arr, idx: arr.numpy()[idx]
|
|
raise RuntimeError("unknown name")
|