mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
849d9c343c
remove unsupported tests Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com>
134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
"""
|
|
Local implementation of our Python venv testing.
|
|
All the bootstrapping is done with our original core work.
|
|
"""
|
|
import copy
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import textwrap as tw
|
|
|
|
from defs.runner_interface import PythonRunnerInterface
|
|
from defs.trt_test_alternative import check_call, check_output, exists, makedirs
|
|
|
|
|
|
class PythonVenvRunnerImpl(PythonRunnerInterface):
|
|
"""
|
|
Object that understands how to run Python scripts in a virtual environment.
|
|
Local implementation of our runner.
|
|
|
|
Args:
|
|
pip_opts (list): Options to pass to pip when installing packages
|
|
venv_dir (str): Path to the virtualenv root directory, or None if this is
|
|
an externally-built virtualenv
|
|
venv_bin (str): Path to the Python executable to use when running tests
|
|
workspace (str): Path to the TURTLE workspace
|
|
"""
|
|
|
|
def __init__(self, pip_opts, venv_dir, venv_bin, workspace):
|
|
self._venv_bin = venv_bin
|
|
self.venv_dir = venv_dir
|
|
self.pip_opts = pip_opts
|
|
self.venv_exe_name = os.path.basename(venv_bin)
|
|
self.venv_exe_dir = os.path.dirname(venv_bin)
|
|
self.workspace = workspace
|
|
if not exists(self.workspace):
|
|
try:
|
|
makedirs(self.workspace)
|
|
except Exception as e:
|
|
print(f"Error creating workspace directory: {e}")
|
|
self._new_env = os.environ.copy()
|
|
|
|
def get_working_directory(self, translate_wsl_path=True):
|
|
"""
|
|
Common interface required by RunnerInterface.
|
|
Both TRTExecutableRunner and VirtualenvRunner both have a workspace attribute but they are
|
|
of different types (function vs variable), as a result, this function is introduced instead.
|
|
"""
|
|
return self.workspace
|
|
|
|
def _get_repro_envvars(self):
|
|
return sorted(os.environ)
|
|
|
|
def set_working_directory(self, value):
|
|
"""
|
|
Common interface required by RunnerInterface.
|
|
Both TRTExecutableRunner and VirtualenvRunner both have a workspace attribute but they are
|
|
of different types (function vs variable), as a result, this function is introduced instead.
|
|
"""
|
|
self.workspace = value
|
|
|
|
def _run_internal(self, script, caller, print_script: bool):
|
|
|
|
if isinstance(script, bytes):
|
|
script = script.decode()
|
|
if print_script:
|
|
print("Run with Python: {}".format(self._venv_bin))
|
|
print("=== BEGIN SCRIPT ===")
|
|
print(script)
|
|
print("=== END SCRIPT =====")
|
|
|
|
if not os.path.exists(self.workspace):
|
|
os.makedirs(self.workspace, exist_ok=True)
|
|
f = tempfile.NamedTemporaryFile(dir=self.workspace,
|
|
mode="w",
|
|
delete=False)
|
|
f.write(script)
|
|
# On Windows, we have to first close the file before we can read it
|
|
f.close()
|
|
try:
|
|
path = f.name
|
|
out = caller([self._venv_bin, path])
|
|
finally:
|
|
os.remove(f.name)
|
|
|
|
return out
|
|
|
|
def run(self, body, caller=check_call, print_script=True):
|
|
"""Run Python code in a script with a predefined prolog."""
|
|
script = ""
|
|
script += tw.dedent(body)
|
|
return self._run_internal(script, caller, print_script)
|
|
|
|
def run_output(self, body):
|
|
"""Runs Python code and captures the output in a string."""
|
|
return self.run(body, caller=check_output, print_script=False)
|
|
|
|
def run_raw(self, script, caller=check_call, print_script=True):
|
|
"""Run Python code without any pre-processing."""
|
|
return self._run_internal(script, caller, print_script)
|
|
|
|
def run_cmd(self,
|
|
args,
|
|
caller=check_call,
|
|
env=None,
|
|
print_script=True,
|
|
**kwargs):
|
|
"""Call <python-exe> <args> on the command-line (can be used to run Python script files)."""
|
|
if not os.path.exists(self.workspace):
|
|
os.makedirs(self.workspace, exist_ok=True)
|
|
|
|
call_args = [self._venv_bin] + args
|
|
if env:
|
|
new_env = copy.deepcopy(os.environ)
|
|
new_env.update(env)
|
|
else:
|
|
new_env = os.environ
|
|
|
|
if caller.__name__ == 'check_output':
|
|
result = subprocess.run(call_args,
|
|
env=new_env,
|
|
check=True,
|
|
capture_output=True,
|
|
**kwargs)
|
|
return result.stdout.decode('utf-8')
|
|
else:
|
|
print(f"Start subprocess with {caller}({call_args}, env={new_env})")
|
|
return caller(call_args, env=new_env, **kwargs)
|
|
|
|
def install_packages(self, packages):
|
|
"""Install Python packages by name."""
|
|
|
|
def install_from_requirements(self, requirements_file):
|
|
pass
|