[XPU] avoid using xpu-smi to query device info (#72136)

* [XPU] avoid using xpu-smi to query device info

* refine code

* return MB when querying memory

* fix bugs

* add unittest

Author: lj970926

paddle/fluid/pybind/pybind.cc

@@ -2940,6 +2940,10 @@ All parameter, weight, gradient are variables in Paddle.
 #ifdef PADDLE_WITH_XPU
   m.def("get_xpu_device_count", platform::GetXPUDeviceCount);
   m.def("xpu_empty_cache", platform::EmptyCache);
+  m.def("get_xpu_device_utilization_rate",
+        platform::GetXPUDeviceUtilizationRate);
+  m.def("get_xpu_device_total_memory", platform::GetXPUDeviceTotalMemory);
+  m.def("get_xpu_device_used_memory", platform::GetXPUDeviceUsedMemory);
 #endif
 
   py::enum_<platform::TracerOption>(m, "TracerOption", py::arithmetic())

paddle/phi/backends/xpu/enforce_xpu.h

@@ -17,6 +17,7 @@ limitations under the License. */
 #ifdef PADDLE_WITH_XPU
 #include <cuda.h>
 #include <cuda_runtime.h>
+#include <xpu/xpuml.h>
 #endif
 
 #include "paddle/phi/backends/xpu/xpu_header.h"
@@ -128,6 +129,16 @@ DEFINE_EXTERNAL_API_TYPE(cudaError_t, cudaSuccess);
     }                                                               \
   } while (0)
 
+// TODO(lijin23): support fine-grained error msg.
+#define PADDLE_ENFORCE_XPUML_SUCCESS(COND)                              \
+  do {                                                                  \
+    auto __cond__ = (COND);                                             \
+    PADDLE_ENFORCE_EQ(                                                  \
+        __cond__,                                                       \
+        XPUML_SUCCESS,                                                  \
+        common::errors::Fatal("XPUML Error, error_code=%d", __cond__)); \
+  } while (0)
+
 }  // namespace xpu
 }  // namespace backends
 }  // namespace phi

paddle/phi/backends/xpu/xpu_info.cc

@@ -50,6 +50,8 @@ class XPUContext;
 namespace backends {
 namespace xpu {
 
+static std::once_flag xpuml_init_flag;
+
 /**************************** Version Management **************************/
 
 //! Get the version of XPU Driver
@@ -221,6 +223,50 @@ void MemcpySyncD2D(void* dst,
 
 /**************************** Others **************************/
 
+int GetXPUDeviceUtilizationRate(int dev_id) {
+  std::call_once(xpuml_init_flag, xpumlInit);
+  if (dev_id == -1) {
+    dev_id = GetXPUCurrentDeviceId();
+  }
+  xpumlDevice_t dev_handle;
+  PADDLE_ENFORCE_XPUML_SUCCESS(
+      xpumlDeviceGetHandleByIndex(dev_id, &dev_handle));
+  xpumlUtilization_t dev_util;
+  PADDLE_ENFORCE_XPUML_SUCCESS(
+      xpumlDeviceGetUtilizationRates(dev_handle, &dev_util));
+  return dev_util.xpu;
+}
+
+int GetXPUDeviceTotalMemory(int dev_id) {
+  std::call_once(xpuml_init_flag, xpumlInit);
+  if (dev_id == -1) {
+    dev_id = GetXPUCurrentDeviceId();
+  }
+
+  xpumlDevice_t dev_handle;
+  PADDLE_ENFORCE_XPUML_SUCCESS(
+      xpumlDeviceGetHandleByIndex(dev_id, &dev_handle));
+  xpumlMemory_t dev_mem_info;
+  PADDLE_ENFORCE_XPUML_SUCCESS(
+      xpumlDeviceGetMemoryInfo(dev_handle, &dev_mem_info));
+  return dev_mem_info.totalGlobalMemory / 1024 / 1024;  // MB
+}
+
+int GetXPUDeviceUsedMemory(int dev_id) {
+  std::call_once(xpuml_init_flag, xpumlInit);
+  if (dev_id == -1) {
+    dev_id = GetXPUCurrentDeviceId();
+  }
+
+  xpumlDevice_t dev_handle;
+  PADDLE_ENFORCE_XPUML_SUCCESS(
+      xpumlDeviceGetHandleByIndex(dev_id, &dev_handle));
+  xpumlMemory_t dev_mem_info;
+  PADDLE_ENFORCE_XPUML_SUCCESS(
+      xpumlDeviceGetMemoryInfo(dev_handle, &dev_mem_info));
+  return dev_mem_info.usedGlobalMemory / 1024 / 1024;  // MB
+}
+
 XPUVersion get_xpu_version(int dev_id) {
   if (dev_id == -1) {
     dev_id = GetXPUCurrentDeviceId();

paddle/phi/backends/xpu/xpu_info.h

@@ -114,6 +114,9 @@ XPUVersion get_xpu_version(int dev_id);
 void set_xpu_debug_level(int level);
 
 int get_xpu_max_ptr_size(int dev_id);
+int GetXPUDeviceUtilizationRate(int dev_id);
+int GetXPUDeviceTotalMemory(int dev_id);
+int GetXPUDeviceUsedMemory(int dev_id);
 
 }  // namespace xpu
 }  // namespace backends

paddle/phi/core/platform/device/xpu/xpu_info.cc

@@ -119,6 +119,18 @@ void EmptyCache() {
   }
 }
 
+int GetXPUDeviceUtilizationRate(int dev_id) {
+  return phi::backends::xpu::GetXPUDeviceUtilizationRate(dev_id);
+}
+
+int GetXPUDeviceTotalMemory(int dev_id) {
+  return phi::backends::xpu::GetXPUDeviceTotalMemory(dev_id);
+}
+
+int GetXPUDeviceUsedMemory(int dev_id) {
+  return phi::backends::xpu::GetXPUDeviceUsedMemory(dev_id);
+}
+
 class RecordedXPUMallocHelper {
  private:
   explicit RecordedXPUMallocHelper(int dev_id, uint64_t limit_size = 0)

paddle/phi/core/platform/device/xpu/xpu_info.h

@@ -94,6 +94,10 @@ bool IsXPUMallocRecorded(int dev_id);
 
 void EmptyCache(void);
 
+int GetXPUDeviceUtilizationRate(int dev_id);
+int GetXPUDeviceTotalMemory(int dev_id);
+int GetXPUDeviceUsedMemory(int dev_id);
+
 }  // namespace platform
 }  // namespace paddle
 #endif

python/paddle/distributed/launch/utils/nvsmi.py

@@ -20,6 +20,7 @@
 import time
 
 import paddle
+from paddle.base import core
 
 
 class Info:
@@ -150,40 +151,23 @@ def query_npu_smi(query=None, index=None, dtype=None):
 
 
 def query_xpu_smi(query=None, index=None, dtype=None):
-    if not has_xpu_smi():
-        return []
-
-    cmd = ["xpu-smi"]
-
+    ret = []
     if not isinstance(dtype, list) or len(dtype) != len(query):
         dtype = [str] * len(query)
 
-    output = subprocess.check_output(cmd, timeout=3)
-    lines = output.decode("utf-8").split(os.linesep)
-    ret = []
-    i = 0
-
-    for line in lines:
-        if not line:
-            continue
-        result = re.split(r',|/|\s+|\|', line)
-        length = len(result)
-        if length not in [23] or "XPU" in result:
-            continue
-        result = [item for item in result if item]
-        info = Info()
-        utilization_xpu = float(re.findall(r'\d+\.\d+|\d+', result[9])[0])
-        mem_total = float(re.findall(r'\d+\.\d+|\d+', result[8])[0])
-        mem_used = float(re.findall(r'\d+\.\d+|\d+', result[7])[0])
+    for dev_id in range(core.get_xpu_device_count()):
+        utilization_xpu = core.get_xpu_device_utilization_rate(dev_id)
+        mem_total = core.get_xpu_device_total_memory(dev_id)
+        mem_used = core.get_xpu_device_used_memory(dev_id)
         result = [
-            i,
+            dev_id,
             utilization_xpu,
             mem_total,
             mem_used,
             (mem_total - mem_used),
             time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
         ]
-        i += 1
+        info = Info()
         for k, v, d in zip(query, result, dtype):
             setattr(info, k.replace(".", "_"), d(v))
         ret.append(info)

test/xpu/test_query_xpu_device_info.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 unittest
+
+import paddle
+from paddle.device import core
+
+
+class TestQueryXPUDeviceInfo(unittest.TestCase):
+    def test_dygraph(self):
+        if core.is_compiled_with_xpu():
+            paddle.disable_static()
+            dev_num = core.get_xpu_device_count()
+            self.assertGreater(
+                dev_num,
+                0,
+                "The environment you run this test does not have any xpu device.",
+            )
+            for dev_id in range(dev_num):
+                self.assertGreaterEqual(
+                    core.get_xpu_device_utilization_rate(dev_id), 0
+                )
+                self.assertGreater(core.get_xpu_device_total_memory(dev_id), 0)
+                self.assertGreaterEqual(
+                    core.get_xpu_device_used_memory(dev_id), 0
+                )
+            paddle.enable_static()
+
+
+if __name__ == "__main__":
+    unittest.main()