From 21e0ea00e28bb5e8a63e740330de3ab5762940a3 Mon Sep 17 00:00:00 2001
From: Xuehai Pan <XuehaiPan@pku.edu.cn>
Date: Mon, 10 Apr 2023 12:55:21 +0000
Subject: [PATCH] fix(api/device): further isolate the `CUDA_VISIBLE_DEVICE`
 parser in a subprocess

---
 nvitop/api/device.py | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/nvitop/api/device.py b/nvitop/api/device.py
index d137cc9..db49158 100644
--- a/nvitop/api/device.py
+++ b/nvitop/api/device.py
@@ -108,6 +108,9 @@ import functools
 import multiprocessing as mp
 import os
 import re
+import subprocess
+import sys
+import textwrap
 import threading
 from collections import OrderedDict
 from typing import Any, Callable, Iterable, NamedTuple
@@ -2488,8 +2491,32 @@ def _parse_cuda_visible_devices(  # pylint: disable=too-many-branches,too-many-s
     gpu_uuids = set(physical_device_attrs)
 
     try:
-        raw_uuids = _parse_cuda_visible_devices_to_uuids(cuda_visible_devices, verbose=False)
-    except libcuda.CUDAError:
+        raw_uuids = (
+            subprocess.check_output(
+                [
+                    sys.executable,
+                    '-c',
+                    textwrap.dedent(
+                        f"""
+                        import nvitop.api.device
+
+                        print(
+                            ','.join(
+                                nvitop.api.device._parse_cuda_visible_devices_to_uuids(
+                                    {cuda_visible_devices!r},
+                                    verbose=False,
+                                ),
+                            ),
+                        )
+                        """,
+                    ),
+                ],
+            )
+            .decode('utf-8')
+            .strip()
+            .split(',')
+        )
+    except subprocess.CalledProcessError:
         pass
     else:
         uuids = [