rename module

Signed-off-by: Xavier Dupre <[email protected]>

onnxscript/tools/benchmark/export_model.py

@@ -60,7 +60,7 @@ def main(args=None):
 
     # Import is delayed so that help is being display faster (without having to import heavy packages).
     import onnxscript.tools
-    import onnxscript.tools.memory_peak as mpk
+    import onnxscript.tools.memory_peak
     import onnxscript.tools.transformers_models
 
     print(
@@ -127,7 +127,7 @@ def main(args=None):
         filename = f"em_{name}.onnx"
 
         memory_session = (
-            mpk.start_spying_on(cuda=kwargs["device"] == "cuda")
+            onnxscript.tools.memory_peak.start_spying_on(cuda=kwargs["device"] == "cuda")
             if kwargs["memory_peak"]
             else None
         )
@@ -148,7 +148,9 @@ def main(args=None):
         if memory_session is not None:
             memory_results = memory_session.stop()
             print(f"[export_model] ends memory monitoring {memory_results}")
-            memory_stats = mpk.flatten(memory_results, prefix="memory_")
+            memory_stats = onnxscript.tools.memory_peak.flatten(
+                memory_results, prefix="memory_"
+            )
         else:
             memory_stats = {}
 

onnxscript/tools/memory_peak_test.py

@@ -7,16 +7,16 @@
 import numpy as np
 import torch
 
-import onnxscript.tools.memory_peak as mpeak
+import onnxscript.tools.memory_peak
 
 
 class TestMemoryPeak(unittest.TestCase):
     def test_memory(self):
-        mem = mpeak.get_memory_rss(os.getpid())
+        mem = onnxscript.tools.memory_peak.get_memory_rss(os.getpid())
         self.assertIsInstance(mem, int)
 
     def test_spy(self):
-        p = mpeak.start_spying_on()
+        p = onnxscript.tools.memory_peak.start_spying_on()
         res = []
         for i in range(10):
             time.sleep(0.005)
@@ -31,7 +31,7 @@ def test_spy(self):
 
     @unittest.skipIf(not torch.cuda.is_available(), reason="CUDA not here")
     def test_spy_cuda(self):
-        p = mpeak.start_spying_on(cuda=True)
+        p = onnxscript.tools.memory_peak.start_spying_on(cuda=True)
         res = []
         for i in range(10):
             time.sleep(0.005)