diff --git a/tests/integration/defs/agg_unit_mem_df.csv b/tests/integration/defs/agg_unit_mem_df.csv index 2629c91b541..a33d05450fd 100644 --- a/tests/integration/defs/agg_unit_mem_df.csv +++ b/tests/integration/defs/agg_unit_mem_df.csv @@ -104,14 +104,14 @@ unittest/trt/model/test_mamba.py,NVIDIA H100,10, unittest/_torch/attention,NVIDIA Graphics Device,4,B200 Bring Up Board unittest/_torch/misc,NVIDIA Graphics Device,4,B200 Bring Up Board unittest/_torch/speculative,NVIDIA Graphics Device,4,B200 Bring Up Board -unittest/_torch/thop/parallel,NVIDIA Graphics Device,4,B200 Bring Up Board +unittest/_torch/thop/parallel,NVIDIA Graphics Device,16,B200 Bring Up Board "unittest/_torch/auto_deploy/unit/singlegpu -k ""not test_trtllm_bench_backend_comparison""",NVIDIA Graphics Device,4,B200 Bring Up Board unittest/_torch/attention,NVIDIA B200,4, unittest/_torch/misc,NVIDIA B200,4, unittest/_torch/speculative,NVIDIA B200,4, -unittest/_torch/thop/parallel,NVIDIA B200,4, +unittest/_torch/thop/parallel,NVIDIA B200,16, "unittest/_torch/auto_deploy/unit/singlegpu -k ""not test_trtllm_bench_backend_comparison""",NVIDIA B200,4, unittest/_torch/attention,NVIDIA H100,4, unittest/_torch/misc,NVIDIA H100,4, unittest/_torch/speculative,NVIDIA H100,2, -unittest/_torch/thop/parallel,NVIDIA H100,4, +unittest/_torch/thop/parallel,NVIDIA H100,16, diff --git a/tests/unittest/_torch/thop/parallel/conftest.py b/tests/unittest/_torch/thop/parallel/conftest.py new file mode 100644 index 00000000000..27911b32157 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/conftest.py @@ -0,0 +1,28 @@ +import os + +import pytest + + +@pytest.fixture(autouse=True, scope='function') +def set_torchinductor_compile_threads(): + """ + Fixture to set TORCHINDUCTOR_COMPILE_THREADS=1 for tests in this directory. + """ + # --- Setup Phase --- + # Save the original value if it exists + original_value = os.environ.get('TORCHINDUCTOR_COMPILE_THREADS') + + # Set the desired value for the test + os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1' + + # Let the test run with the new environment variable + yield + + # --- Teardown Phase --- + # Restore the original environment state after the test is done + if original_value is None: + # If the variable didn't exist before, remove it + del os.environ['TORCHINDUCTOR_COMPILE_THREADS'] + else: + # Otherwise, restore its original value + os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = original_value diff --git a/tests/unittest/pytest.ini b/tests/unittest/pytest.ini index f1d72c7dbe6..1e93eae77ab 100644 --- a/tests/unittest/pytest.ini +++ b/tests/unittest/pytest.ini @@ -1,4 +1,5 @@ [pytest] +xdist_start_method = spawn asyncio_default_fixture_loop_scope = module threadleak = True threadleak_exclude = asyncio_\d+