Add device auto-discovery and cli option

Add `--device` cli option to specify which CI device is
being tested.  If the cli option is not specified, then
we fallback to the usual GAUDI2_CI environment variable.
Finally, if neither is specified, then try to get the
device name automatically from the habana frameworks
API.

This setting is to prepare for gaudi3 specific (and
future devices) test cases and baseline references.

Eventually, we can deprecate the GAUDI2_CI environment
variable once the rest of the test framework is
migrated to the new device declaration.

Signed-off-by: U. Artie Eoff <[email protected]>

conftest.py

@@ -1,9 +1,12 @@
 import json
 import logging
+import os
 from pathlib import Path
 
 import pytest
 
+import tests.utils as oh_testutils
+
 
 BASELINE_DIRECTORY = Path(__file__).parent.resolve() / Path("tests") / Path("baselines") / Path("fixture")
 
@@ -89,6 +92,7 @@ def __str___(self):
 def pytest_addoption(parser):
     parser.addoption("--token", action="store", default=None)
     parser.addoption("--rebase", action="store_true", help="rebase baseline references from current run")
+    parser.addoption("--device", action="store", default=None)
 
 
 @pytest.fixture
@@ -99,6 +103,30 @@ def token(request):
 def pytest_sessionstart(session):
     session.stash["baseline"] = Baseline(session)
 
+    # User command-line option takes highest priority
+    if session.config.option.device is not None:
+        device = str(session.config.option.device).lower()
+    # User GAUDI2_CI environment variable takes second priority for backwards compatibility
+    elif "GAUDI2_CI" in os.environ:
+        device = "gaudi2" if os.environ["GAUDI2_CI"] == "1" else "gaudi1"
+    # Try to automatically detect it
+    else:
+        import habana_frameworks.torch.hpu as torch_hpu
+
+        name = torch_hpu.get_device_name().strip()
+        device = "unknown" if not name else name.split()[-1].lower()
+
+    # torch_hpu.get_device_name() returns GAUDI for G1
+    if "gaudi" == device:
+        # use "gaudi1" since this is used in tests, baselines, etc.
+        device = "gaudi1"
+
+    oh_testutils.OH_DEVICE_CONTEXT = device
+
+
+def pytest_report_header():
+    return [f"device context: {oh_testutils.OH_DEVICE_CONTEXT}"]
+
 
 def pytest_sessionfinish(session):
     session.stash["baseline"].finalize()

tests/utils.py

@@ -103,3 +103,6 @@
 MODELS_TO_TEST_FOR_SPEECH_RECOGNITION = ["wav2vec2", "whisper"]
 
 MODELS_TO_TEST_FOR_IMAGE_TEXT = ["clip"]
+
+# This will be configured by conftest.py at the start of the pytest session
+OH_DEVICE_CONTEXT = "unknown"