Fix inference unit test import error catching

.github/workflows/nv-inference.yml

@@ -60,4 +60,4 @@ jobs:
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
           cd tests
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 -m 'inference' unit/
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'inference' unit/ --torch_ver="1.8" --cuda_ver="11.1"

.github/workflows/nv-nightly.yml

@@ -41,12 +41,12 @@ jobs:
       - name: Install deepspeed
         run: |
           pip uninstall --yes deepspeed
-          pip install .[dev,1bit,autotuning,sparse_attn]
+          pip install .[dev,1bit,autotuning,sparse_attn,inf]
           ds_report
 
       - name: Unit tests
         run: |
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
           cd tests
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'nightly' unit/
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'nightly' unit/ --torch_ver="1.8" --cuda_ver="11.1"

.github/workflows/nv-torch12-p40.yml

@@ -59,4 +59,4 @@ jobs:
         run: |
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
           cd tests
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4 unit/ --torch_ver="1.2" --cuda_ver="10"

.github/workflows/nv-torch18-v100.yml

@@ -60,5 +60,5 @@ jobs:
           unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
           if [[ -d ./torch-extensions ]]; then rm -rf ./torch-extensions; fi
           cd tests
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4  unit/
-          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -n 4  unit/ --torch_ver="1.8" --cuda_ver="11.1"
+          TORCH_EXTENSIONS_DIR=./torch-extensions pytest --color=yes --durations=0 --forked --verbose -m 'sequential' unit/ --torch_ver="1.8" --cuda_ver="11.1"

tests/conftest.py

@@ -1,9 +1,45 @@
 # tests directory-specific settings - this file is run automatically by pytest before any tests are run
 
 import sys
+import pytest
 from os.path import abspath, dirname, join
+import torch
+import warnings
 
 # allow having multiple repository checkouts and not needing to remember to rerun
 # 'pip install -e .[dev]' when switching between checkouts and running tests.
 git_repo_path = abspath(join(dirname(dirname(__file__)), "src"))
 sys.path.insert(1, git_repo_path)
+
+
+def pytest_addoption(parser):
+    parser.addoption("--torch_ver", default=None, type=str)
+    parser.addoption("--cuda_ver", default=None, type=str)
+
+
+def validate_version(expected, found):
+    version_depth = expected.count('.') + 1
+    found = '.'.join(found.split('.')[:version_depth])
+    return found == expected
+
+
+@pytest.fixture(scope="session", autouse=True)
+def check_environment(pytestconfig):
+    expected_torch_version = pytestconfig.getoption("torch_ver")
+    expected_cuda_version = pytestconfig.getoption("cuda_ver")
+    if expected_torch_version is None:
+        warnings.warn(
+            "Running test without verifying torch version, please provide an expected torch version with --torch_ver"
+        )
+    elif not validate_version(expected_torch_version, torch.__version__):
+        pytest.exit(
+            f"expected torch version {expected_torch_version} did not match found torch version {torch.__version__}",
+            returncode=2)
+    if expected_cuda_version is None:
+        warnings.warn(
+            "Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver"
+        )
+    elif not validate_version(expected_cuda_version, torch.version.cuda):
+        pytest.exit(
+            f"expected cuda version {expected_cuda_version} did not match found cuda version {torch.version.cuda}",
+            returncode=2)

tests/unit/test_inference.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import time
 import torch
 import pytest
@@ -9,17 +10,18 @@
 from .common import distributed_test
 from packaging import version as pkg_version
 from deepspeed.ops.op_builder import OpBuilder
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import HfApi
 
-try:
+
+@pytest.fixture(scope="module", autouse=True)
+def lm_eval_imports():
+    global lm_eval
     import lm_eval
     import lm_eval.models
     import lm_eval.tasks
-    from lm_eval.evaluator import evaluate
-    from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-    from huggingface_hub import HfApi
-except ImportError:
-    pytest.skip("please install w. [inf] extra to run this test",
-                allow_module_level=True)
+    import lm_eval.evaluator
+
 
 rocm_version = OpBuilder.installed_rocm_version()
 if rocm_version != (0, 0):
@@ -295,7 +297,7 @@ def _go():
 
         torch.cuda.synchronize()
         start = time.time()
-        bs_output = evaluate(lm=lm, task_dict=task_dict)
+        bs_output = lm_eval.evaluator.evaluate(lm=lm, task_dict=task_dict)
         torch.cuda.synchronize()
         bs_time = time.time() - start
 
@@ -311,7 +313,7 @@ def _go():
         setattr(lm, model_family, ds_model)
         torch.cuda.synchronize()
         start = time.time()
-        ds_output = evaluate(lm=lm, task_dict=task_dict)
+        ds_output = lm_eval.evaluator.evaluate(lm=lm, task_dict=task_dict)
         torch.cuda.synchronize()
         ds_time = time.time() - start