cpu only env compat

Signed-off-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com>

Author: tongyuantongyu

jenkins/L0_Test.groovy

@@ -1078,7 +1078,7 @@ def getPytestBaseCommandLine(
     if (stageName.contains("-Ray-")) {
         testCmdLine += ["--run-ray"]
     }
-    def unittestMarkExpr = (stageName.startsWith("CPU-") || stageName.contains("-CPU-")) ? "cpu_only and not disabled" : "not cpu_only"
+    def unittestMarkExpr = (stageName.startsWith("CPU-")) ? "cpu_only and not disabled" : "not cpu_only"
     testCmdLine += ["--unittest-markexpr=${unittestMarkExpr}"]
     if (extraArgs) {
         testCmdLine += extraArgs
@@ -2911,7 +2911,7 @@ def renderTestDB(pipeline, testContext, llmSrc, stageName, preDefinedMakoOpts=nu
 
     if (!makoOpts) {
         def makoArgs = getMakoArgsFromStageName(stageName)
-        if (stageName.startsWith("CPU-") || stageName.contains("-CPU-")) {
+        if (stageName.startsWith("CPU-")) {
             def cpuName = env.targetArch == AARCH64_TRIPLE ? "aarch64" : "x86_64"
             makoOpts = transformMakoArgsToJson(
                 ["Mako options:"] + makoArgs + [
@@ -3518,7 +3518,10 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
         sh "nproc && free -g && hostname"
         echoNodeAndGpuInfo(pipeline, stageName)
         sh "cat ${MODEL_CACHE_DIR}/README"
-        sh "nvidia-smi && nvidia-smi -q && nvidia-smi topo -m"
+        sh "nvidia-smi && nvidia-smi -q && nvidia-smi topo -m || echo nvidia-smi missing"
+        if (stageName.startsWith("CPU-")) {
+            sh "ln -s /usr/local/cuda/compat/lib.real /usr/local/cuda/compat/lib"
+        }
         sh "df -h"
 
         // setup HF_HOME to cache model and datasets

tensorrt_llm/profiler.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import contextlib
 import time
 from functools import partial
 from typing import Literal, Optional, Tuple, Union
@@ -110,18 +111,23 @@ def summary():
 MemUnitType = Literal["GiB", "MiB", "KiB"]
 
 
-class PyNVMLContext:
-    def __enter__(self):
-        if pynvml is not None:
+@contextlib.contextmanager
+def pynvml_context():
+    has_pynvml = pynvml is not None
+    if has_pynvml:
+        try:
             pynvml.nvmlInit()
+        except pynvml.NVMLError:
+            has_pynvml = False
 
-    def __exit__(self, type, value, traceback):
-        if pynvml is not None:
-            pynvml.nvmlShutdown()
+    yield
+
+    if has_pynvml:
+        pynvml.nvmlShutdown()
 
 
 if pynvml is not None:
-    with PyNVMLContext():
+    with pynvml_context():
         _device_get_memory_info_fn = partial(
             pynvml.nvmlDeviceGetMemoryInfo,
             version=pynvml.nvmlMemory_v2,
@@ -147,7 +153,7 @@ def device_memory_info(device: Optional[Union[torch.device, int]] = None) -> Tup
         if device is None:
             device = torch.cuda.current_device()
         index = device.index if isinstance(device, torch.device) else device
-        with PyNVMLContext():
+        with pynvml_context():
             handle = pynvml.nvmlDeviceGetHandleByIndex(index)
             mem_info = _device_get_memory_info_fn(handle)
         return mem_info.used, mem_info.free, mem_info.total

tests/integration/defs/conftest.py

@@ -1797,13 +1797,17 @@ def get_gpu_device_list():
         suffix = ".exe" if is_windows() else ""
         # TODO: Use NRSU because we can't assume nvidia-smi across all platforms.
         cmd = " ".join(["nvidia-smi" + suffix, "-L"])
-        output = check_output(cmd, shell=True, cwd=temp_dirname)
+        try:
+            output = check_output(cmd, shell=True, cwd=temp_dirname)
+        except sp.CalledProcessError:
+            return []
     return [l.strip() for l in output.strip().split("\n")]
 
 
 def check_device_contain(keyword_list):
     "check device not contain keyword"
-    device = get_gpu_device_list()[0]
+    devices = get_gpu_device_list()
+    device = devices[0] if devices else ""
     return any(keyword in device for keyword in keyword_list)