feat(mge/device): add python binding for get_mem_status_bytes

GitOrigin-RevId: aa0d3264d6042b1f3c7dfe243ef9a97794868c42

imperative/python/megengine/device.py

@@ -8,6 +8,7 @@
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 import os
 import re
+from typing import Optional
 
 from .core._imperative_rt.common import CompNode, DeviceType
 from .core._imperative_rt.common import set_prealloc_config as _set_prealloc_config
@@ -17,6 +18,7 @@ __all__ = [
     "get_device_count",
     "get_default_device",
     "set_default_device",
+    "get_mem_status_bytes",
     "set_prealloc_config",
     "DeviceType",
 ]
@@ -92,6 +94,16 @@ def get_default_device() -> str:
     return CompNode._get_default_device()
 
 
+def get_mem_status_bytes(device: Optional[str] = None):
+    r"""
+    Get total and free memory on the computing device in bytes.
+    """
+    if device is None:
+        device = get_default_device()
+    tot, free = CompNode(device).get_mem_status_bytes
+    return tot, free
+
+
 set_default_device(os.getenv("MGE_DEFAULT_DEVICE", "xpux"))
 
 

imperative/python/src/common.cpp

@@ -58,6 +58,9 @@ void init_common(py::module m) {
         .def_property_readonly("logical_name", [](const CompNode& cn) {
             return cn.to_string_logical();
         })
+        .def_property_readonly("get_mem_status_bytes", [](const CompNode& cn) {
+            return cn.get_mem_status_bytes();
+        })
         .def("create_event", &CompNode::create_event, py::arg("flags") = 0ul)
         .def("_set_default_device", &set_default_device)
         .def("_get_default_device", &get_default_device)