Arm backend: Make pytest collection seeding deterministic

backends/arm/test/conftest.py

@@ -26,6 +26,19 @@ def pytest_configure(config):
         pytest._test_options["llama_inputs"] = config.option.llama_inputs  # type: ignore[attr-defined]
 
     logging.basicConfig(stream=sys.stdout)
+    seed, seed_label = _setup_random_seed()
+    config._test_seed = seed
+    config._test_seed_label = seed_label
+
+    if os.environ.get("TEST_RUNTIME_IS_NOT_OSS", "0") != "1":
+        # This imports/uses torch early, which doesn't work in some Buck2 test environments.
+        # Since this only makes randomness deterministic (reducing flakiness), it's mainly meant for
+        # local/OSS project test runs.
+        _set_random_seed(seed)
+
+
+def pytest_report_header(config):
+    return config._test_seed_label
 
 
 def pytest_collection_modifyitems(config, items):
@@ -63,38 +76,44 @@ def pytest_sessionfinish(session, exitstatus):
 
 
 @pytest.fixture(autouse=True)
-def set_random_seed():
+def set_random_seed(request):
     """Control random numbers in Arm test suite. Default behavior is to use a
     fixed seed (0), which ensures reproducible tests. Use the env variable
-    ARM_TEST_SEED to set a custom seed, or set it to RANDOM for random seed
-    behavior.
+    TEST_SEED to set a custom session seed, or set it to RANDOM to choose a
+    random session seed.
 
     Examples:
     As default use fixed seed (0) for reproducible tests
         pytest --config-file=/dev/null --verbose -s --color=yes  backends/arm/test/ops/test_avg_pool.py -k <TESTCASE>
-    Use a random seed for each test
-        ARM_TEST_SEED=RANDOM pytest --config-file=/dev/null --verbose -s --color=yes  backends/arm/test/ops/test_avg_pool.py -k <TESTCASE>
+    Use a random seed for the test session
+        TEST_SEED=RANDOM pytest --config-file=/dev/null --verbose -s --color=yes  backends/arm/test/ops/test_avg_pool.py -k <TESTCASE>
     Rerun with a specific seed
-        ARM_TEST_SEED=3478246 pytest --config-file=/dev/null --verbose -s --color=yes  backends/arm/test/ops/test_avg_pool.py -k <TESTCASE>
+        TEST_SEED=3478246 pytest --config-file=/dev/null --verbose -s --color=yes  backends/arm/test/ops/test_avg_pool.py -k <TESTCASE>
 
     """
-    import torch
+    _set_random_seed(request.config._test_seed)
+
 
-    seed_env = os.environ.get("ARM_TEST_SEED", "0")
+def _setup_random_seed():
+    seed_env = os.environ.get("TEST_SEED", "0")
     if seed_env == "RANDOM":
         random.seed()  # reset seed, in case any other test has fiddled with it
         seed = random.randint(0, 2**32 - 1)  # nosec B311 - non-crypto seed for tests
-        torch.manual_seed(seed)
+        seed_label = f"TEST_SEED=RANDOM using:{seed}"
     elif str.isdigit(seed_env):
         seed = int(seed_env)
-        random.seed(seed)
-        torch.manual_seed(seed)
+        seed_label = f"TEST_SEED={seed}"
     else:
-        raise TypeError(
-            "ARM_TEST_SEED env variable must be integers or the string RANDOM"
-        )
+        raise TypeError("TEST_SEED env variable must be integers or the string RANDOM")
+
+    return seed, seed_label
+
+
+def _set_random_seed(seed):
+    import torch
 
-    print(f" ARM_TEST_SEED={seed} ", end=" ")
+    random.seed(seed)
+    torch.manual_seed(seed)
 
 
 # ==== End of Pytest fixtures =====

backends/arm/test/targets.bzl

@@ -96,6 +96,19 @@ def define_arm_tests():
     for test_file in test_files:
         test_file_name = paths.basename(test_file)
         test_name = test_file_name.replace("test_", "").replace(".py", "")
+        test_env = {
+            "TEST_RUNTIME_IS_NOT_OSS": "1" if not runtime.is_oss else "0",
+        }
+        if not runtime.is_oss and _ENABLE_VGF:
+            test_env.update({
+                "MODEL_CONVERTER_PATH": "$(location fbsource//third-party/pypi/ai-ml-sdk-model-converter/0.9.0:model-converter-bin)",
+                "MODEL_CONVERTER_LIB_DIR": "$(location fbsource//third-party/nvidia-nsight-systems:linux-x86_64)/host-linux-x64",
+                "LAVAPIPE_LIB_PATH": "$(location fbsource//third-party/mesa:vulkan_lvp)",
+                "EMULATION_LAYER_TENSOR_SO": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:libVkLayer_Tensor)",
+                "EMULATION_LAYER_GRAPH_SO": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:libVkLayer_Graph)",
+                "EMULATION_LAYER_TENSOR_JSON": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:VkLayer_Tensor_json)",
+                "EMULATION_LAYER_GRAPH_JSON": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:VkLayer_Graph_json)",
+            })
 
         python_pytest(
             name = test_name,
@@ -105,15 +118,7 @@ def define_arm_tests():
             compile = "with-source",
             typing = False,
             skip_on_mode_mac = True,
-            env = {} if runtime.is_oss else ({
-                "MODEL_CONVERTER_PATH": "$(location fbsource//third-party/pypi/ai-ml-sdk-model-converter/0.9.0:model-converter-bin)",
-                "MODEL_CONVERTER_LIB_DIR": "$(location fbsource//third-party/nvidia-nsight-systems:linux-x86_64)/host-linux-x64",
-                "LAVAPIPE_LIB_PATH": "$(location fbsource//third-party/mesa:vulkan_lvp)",
-                "EMULATION_LAYER_TENSOR_SO": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:libVkLayer_Tensor)",
-                "EMULATION_LAYER_GRAPH_SO": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:libVkLayer_Graph)",
-                "EMULATION_LAYER_TENSOR_JSON": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:VkLayer_Tensor_json)",
-                "EMULATION_LAYER_GRAPH_JSON": "$(location fbsource//third-party/arm-ml-emulation-layer/v0.9.0/src:VkLayer_Graph_json)",
-            } if _ENABLE_VGF else {}),
+            env = test_env,
             preload_deps = [
                 "//executorch/kernels/quantized:custom_ops_generated_lib",
             ] + ([] if runtime.is_oss or not _ENABLE_VGF else [