refactor logic

Author: wwwjn

.github/workflows/integration_test_8gpu_features.yaml

@@ -50,4 +50,4 @@ jobs:
         USE_CPP=0 python -m pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126
 
         mkdir artifacts-to-be-uploaded
-        python -m tests.integration_tests.features artifacts-to-be-uploaded --ngpu 8
+        python -m tests.integration_tests.run_tests --test_suite features artifacts-to-be-uploaded --ngpu 8

.github/workflows/integration_test_8gpu_h100.yaml

@@ -53,4 +53,4 @@ jobs:
         mkdir artifacts-to-be-uploaded
 
         # Enable CPP stacktraces for debugging symmetric memory initialization errors.
-        TORCH_SHOW_CPP_STACKTRACES=1 python -m tests.integration_tests.h100 artifacts-to-be-uploaded --ngpu 8
+        TORCH_SHOW_CPP_STACKTRACES=1 python -m tests.integration_tests.run_tests --test_suite h100 artifacts-to-be-uploaded --ngpu 8

.github/workflows/integration_test_8gpu_models.yaml

@@ -50,4 +50,4 @@ jobs:
         USE_CPP=0 python -m pip install --pre torchao --index-url https://download.pytorch.org/whl/nightly/cu126
 
         mkdir artifacts-to-be-uploaded
-        python -m tests.integration_tests.models artifacts-to-be-uploaded --ngpu 8
+        python -m tests.integration_tests.run_tests --test_suite models artifacts-to-be-uploaded --ngpu 8

tests/integration_tests/ft.py

@@ -9,6 +9,7 @@
 import logging
 import os
 import subprocess
+from typing import List
 
 from .features import OverrideDefinitions
 
@@ -90,14 +91,12 @@ def run_single_test(test_flavor: OverrideDefinitions, full_path: str, output_dir
             )
 
 
-def run_tests(args):
-    integration_tests_flavors = build_test_list()
-
+def run_tests(args, test_list: List[OverrideDefinitions]):
     if args.ngpu < 8:
         logger.info("Skipping TorchFT integration tests as we need 8 GPUs.")
         return
 
-    for test_flavor in integration_tests_flavors:
+    for test_flavor in test_list:
         # Filter by test_name if specified
         if args.test_name != "all" and test_flavor.test_name != args.test_name:
             continue
@@ -138,7 +137,11 @@ def main():
 
     if not os.path.exists(args.output_dir):
         os.makedirs(args.output_dir)
-    run_tests(args)
+    if os.listdir(args.output_dir):
+        raise RuntimeError("Please provide an empty output directory.")
+
+    test_list = build_ft_test_list()
+    run_tests(args, test_list)
 
 
 if __name__ == "__main__":

tests/integration_tests/h100.py

@@ -80,10 +80,11 @@ def build_h100_tests_list():
     return integration_tests_flavors
 
 
-def run_tests(args):
+def run_tests(args, test_list=None):
     """Run all H100 integration tests"""
     # build integration tests list
-    test_list = build_test_list()
+    if test_list is None:
+        test_list = build_h100_tests_list()
 
     for test_flavor in test_list:
         # Filter by test_name if specified
@@ -105,6 +106,8 @@ def run_tests(args):
                 f" because --ngpu arg is {args.ngpu}"
             )
         else:
+            # Import run_single_test from run_tests.py
+            from tests.integration_tests.run_tests import run_single_test
             run_single_test(test_flavor, args.config_path, args.output_dir)
 
 
@@ -128,7 +131,9 @@ def main():
         os.makedirs(args.output_dir)
     if os.listdir(args.output_dir):
         raise RuntimeError("Please provide an empty output directory.")
-    run_tests(args)
+    
+    test_list = build_h100_tests_list()
+    run_tests(args, test_list)
 
 
 if __name__ == "__main__":

tests/integration_tests/run_tests.py

@@ -16,7 +16,6 @@
 
 _TEST_SUITES_FUNCTION = {
     "features": build_features_test_list,
-    "ft": build_ft_test_list,
     "models": build_model_tests_list,
     "h100": build_h100_tests_list,
 }
@@ -99,7 +98,7 @@ def main():
     parser.add_argument(
         "--test_suite",
         default="",
-        choices=["features", "ft", "models", "h100"],
+        choices=["features", "models", "h100"],
         help="Which test suite to run. If not specified, torchtitan composibility tests will be run",
     )
     parser.add_argument(