Prevent hangs in CI during parallel run compilation (#2844)

deepspeedai · Jun 26, 2023 · 2b2be85 · 2b2be85
1 parent 1b88839
commit 2b2be85
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/.github/workflows/nv-torch-latest-cpu.yml b/.github/workflows/nv-torch-latest-cpu.yml
@@ -14,7 +14,7 @@ concurrency:
 
 jobs:
   unit-tests:
-    runs-on: ubuntu-20.04
+    runs-on: [self-hosted, cpu]
 
     steps:
       - uses: actions/checkout@v3

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -39,6 +39,14 @@ def validate_version(expected, found):
     return found == expected
 
 
+# This fixture prevents hangs when 2+ pytest processes try to compile the same
+# code at once and causes deadlocks / hangs
+@pytest.fixture(scope="session", autouse=True)
+def set_torch_ext_dir(worker_id):
+    torch_ext_dir = os.environ["TORCH_EXTENSIONS_DIR"]
+    os.environ["TORCH_EXTENSIONS_DIR"] = os.path.join(torch_ext_dir, worker_id)
+
+
 @pytest.fixture(scope="session", autouse=True)
 def check_environment(pytestconfig):
     expected_torch_version = pytestconfig.getoption("torch_ver")