[TD] Enable TD on AVX related configs (#125482)

On test configs `nogpu_AVX512` and `nogpu_NO_AVX2`, which are the next longest jobs on trunk after windows
Pull Request resolved: https://github.com/pytorch/pytorch/pull/125482
Approved by: https://github.com/huydhn
This commit is contained in:
Catherine Lee
2024-05-06 22:02:16 +00:00
committed by PyTorch MergeBot
parent 8c74162074
commit 1b3fd83ab2

View File

@ -37,6 +37,7 @@ from torch.testing._internal.common_utils import (
TEST_WITH_ASAN,
TEST_WITH_CROSSREF,
TEST_WITH_ROCM,
TEST_WITH_SLOW,
TEST_WITH_SLOW_GRADCHECK,
)
@ -74,7 +75,7 @@ HAVE_TEST_SELECTION_TOOLS = True
# Make sure to remove REPO_ROOT after import is done
sys.path.remove(str(REPO_ROOT))
TEST_CONFIG = os.getenv("TEST_CONFIG", "")
RERUN_DISABLED_TESTS = os.getenv("PYTORCH_TEST_RERUN_DISABLED_TESTS", "0") == "1"
DISTRIBUTED_TEST_PREFIX = "distributed"
INDUCTOR_TEST_PREFIX = "inductor"
@ -491,14 +492,11 @@ def run_test(
and not is_cpp_test
and "-n" not in command
)
is_slow = "slow" in os.environ.get("TEST_CONFIG", "") or "slow" in os.environ.get(
"BUILD_ENVRIONMENT", ""
)
timeout = (
None
if not options.enable_timeout
else THRESHOLD * 6
if is_slow
if TEST_WITH_SLOW
else THRESHOLD * 3
if should_retry
and isinstance(test_module, ShardedTest)
@ -1183,15 +1181,16 @@ def parse_args():
or TEST_WITH_ASAN
or (
strtobool(os.environ.get("TD_DISTRIBUTED", "False"))
and os.getenv("TEST_CONFIG") == "distributed"
and TEST_CONFIG == "distributed"
and TEST_CUDA
)
or (IS_WINDOWS and not TEST_CUDA)
or TEST_CONFIG == "nogpu_AVX512"
or TEST_CONFIG == "nogpu_NO_AVX2"
)
and get_pr_number() is not None
and not strtobool(os.environ.get("NO_TD", "False"))
and "slow" not in os.getenv("TEST_CONFIG", "")
and "slow" not in os.getenv("BUILD_ENVIRONMENT", ""),
and not TEST_WITH_SLOW,
)
parser.add_argument(
"additional_unittest_args",