fix: session start fixture correctly handles being launched by mpirun (…

…#486) Signed-off-by: Terry Kong <[email protected]>
NVIDIA · Jan 21, 2025 · 9512ee8 · 9512ee8
1 parent 735c415
commit 9512ee8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -401,7 +401,9 @@ def pytest_collection_modifyitems(config, items):
 
 def pytest_sessionstart(session):
     # Remove the file at the start of the session, if it exists
-    if os.path.exists(SUCCESS_FILE) and os.environ["LOCAL_RANK"] == "0":
+    if os.path.exists(SUCCESS_FILE) and (
+        os.environ.get("LOCAL_RANK", None) == "0" or os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", None) == "0"
+    ):
         os.remove(SUCCESS_FILE)