diff --git a/noxfile.py b/noxfile.py index eff8136fbb..b17b506b7a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -5,7 +5,7 @@ # Use system installed Python packages PYT_PATH = ( - "/opt/conda/lib/python3.8/site-packages" + "/usr/local/lib/python3.8/dist-packages" if not "PYT_PATH" in os.environ else os.environ["PYT_PATH"] ) @@ -202,6 +202,93 @@ def run_base_tests(session): else: session.run_always("pytest", test) +def run_fx_core_tests(session): + print("Running FX core tests") + session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test")) + tests = [ + "core", + ] + for test in tests: + if USE_HOST_DEPS: + session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH}) + else: + session.run_always("pytest", test) + +def run_fx_converter_tests(session): + print("Running FX converter tests") + session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test")) + tests = [ + "converters", + ] + # Skipping this test as it fails inside NGC container with the following error. + # Error Code 4: Internal Error (Could not find any implementation for node conv due to insufficient workspace. See verbose log for requested sizes.) + skip_tests = "-k not conv3d" + for test in tests: + if USE_HOST_DEPS: + session.run_always("pytest", test, skip_tests, env={"PYTHONPATH": PYT_PATH}) + else: + session.run_always("pytest", test, skip_tests) + +def run_fx_lower_tests(session): + print("Running FX passes and trt_lower tests") + session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test")) + tests = [ + "passes/test_multi_fuse_trt.py", + # "passes/test_fuse_permute_linear_trt.py", + "passes/test_remove_duplicate_output_args.py", + "passes/test_fuse_permute_matmul_trt.py", + #"passes/test_graph_opts.py" + "trt_lower", + ] + for test in tests: + if USE_HOST_DEPS: + session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH}) + else: + session.run_always("pytest", test) + +def run_fx_quant_tests(session): + print("Running FX Quant tests") + session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test")) + tests = [ + "quant", + ] + # Skipping this test as it fails inside NGC container with the following error. + # ImportError: cannot import name 'ObservationType' from 'torch.ao.quantization.backend_config.observation_type' + skip_tests = "-k not conv_add_standalone_module" + for test in tests: + if USE_HOST_DEPS: + session.run_always("pytest", test, skip_tests, env={"PYTHONPATH": PYT_PATH}) + else: + session.run_always("pytest", test, skip_tests) + +def run_fx_tracer_tests(session): + print("Running FX Tracer tests") + session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test")) + # skipping a test since it depends on torchdynamo + # Enable this test once NGC moves to latest pytorch which has dynamo integrated. + tests = [ + "tracer/test_acc_shape_prop.py", + "tracer/test_acc_tracer.py", + #"tracer/test_dispatch_tracer.py" + ] + for test in tests: + if USE_HOST_DEPS: + session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH}) + else: + session.run_always("pytest", test) + +def run_fx_tools_tests(session): + print("Running FX tools tests") + session.chdir(os.path.join(TOP_DIR, "py/torch_tensorrt/fx/test")) + tests = [ + "tools", + ] + for test in tests: + if USE_HOST_DEPS: + session.run_always("pytest", test, env={"PYTHONPATH": PYT_PATH}) + else: + session.run_always("pytest", test) + def run_model_tests(session): print("Running model tests") @@ -309,6 +396,35 @@ def run_l0_api_tests(session): run_base_tests(session) cleanup(session) +def run_l0_fx_tests(session): + if not USE_HOST_DEPS: + install_deps(session) + install_torch_trt(session) + run_fx_core_tests(session) + run_fx_converter_tests(session) + run_fx_lower_tests(session) + cleanup(session) + +def run_l0_fx_core_tests(session): + if not USE_HOST_DEPS: + install_deps(session) + install_torch_trt(session) + run_fx_core_tests(session) + cleanup(session) + +def run_l0_fx_converter_tests(session): + if not USE_HOST_DEPS: + install_deps(session) + install_torch_trt(session) + run_fx_converter_tests(session) + cleanup(session) + +def run_l0_fx_lower_tests(session): + if not USE_HOST_DEPS: + install_deps(session) + install_torch_trt(session) + run_fx_lower_tests(session) + cleanup(session) def run_l0_dla_tests(session): if not USE_HOST_DEPS: @@ -327,7 +443,6 @@ def run_l1_model_tests(session): run_model_tests(session) cleanup(session) - def run_l1_int8_accuracy_tests(session): if not USE_HOST_DEPS: install_deps(session) @@ -337,6 +452,14 @@ def run_l1_int8_accuracy_tests(session): run_int8_accuracy_tests(session) cleanup(session) +def run_l1_fx_tests(session): + if not USE_HOST_DEPS: + install_deps(session) + install_torch_trt(session) + run_fx_quant_tests(session) + run_fx_tracer_tests(session) + run_fx_tools_tests(session) + cleanup(session) def run_l2_trt_compatibility_tests(session): if not USE_HOST_DEPS: @@ -360,6 +483,25 @@ def l0_api_tests(session): """When a developer needs to check correctness for a PR or something""" run_l0_api_tests(session) +@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) +def l0_fx_tests(session): + """When a developer needs to check correctness for a PR or something""" + run_l0_fx_tests(session) + +@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) +def l0_fx_core_tests(session): + """When a developer needs to check correctness for a PR or something""" + run_l0_fx_core_tests(session) + +@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) +def l0_fx_converter_tests(session): + """When a developer needs to check correctness for a PR or something""" + run_l0_fx_converter_tests(session) + +@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) +def l0_fx_lower_tests(session): + """When a developer needs to check correctness for a PR or something""" + run_l0_fx_lower_tests(session) @nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) def l0_dla_tests(session): @@ -372,6 +514,10 @@ def l1_model_tests(session): """When a user needs to test the functionality of standard models compilation and results""" run_l1_model_tests(session) +@nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) +def l1_fx_tests(session): + """When a user needs to test the functionality of standard models compilation and results""" + run_l1_fx_tests(session) @nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) def l1_int8_accuracy_tests(session): @@ -388,4 +534,4 @@ def l2_trt_compatibility_tests(session): @nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True) def l2_multi_gpu_tests(session): """Makes sure that Torch-TensorRT can operate on multi-gpu systems""" - run_l2_multi_gpu_tests(session) + run_l2_multi_gpu_tests(session) \ No newline at end of file