Add file to help benchmark, memory measure (#29)

* bench

* fix test

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.3.0
 +++++
 
+* :pr:`29`: adds helpers to measure the memory peak and run benchmark
+  on different processes
 * :pr:`28`: adds command line to print out the configuration for a model id,
   support image-text-to-text
 * :pr:`26`: creates a folder ``helpers`` to gather all the functions

_doc/api/helpers/bench_run.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.bench_run
+=================================
+
+.. automodule:: onnx_diagnostic.helpers.bench_run
+    :members:
+    :no-undoc-members:

_doc/api/helpers/index.rst

@@ -7,8 +7,10 @@ onnx_diagnostic.helpers
     :caption: submodules
 
     args_helper
+    bench_run
     cache_helper
     helper
+    memory_peak
     onnx_helper
     ort_session
     torch_test_helper

_doc/api/helpers/memory_peak.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.helpers.memory_peak
+===================================
+
+.. automodule:: onnx_diagnostic.helpers.memory_peak
+    :members:
+    :no-undoc-members:

_unittests/ut_helpers/test_bench_run.py

@@ -0,0 +1,166 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
+from onnx_diagnostic.helpers import max_diff
+from onnx_diagnostic.helpers.bench_run import (
+    BenchmarkError,
+    _cmd_line,
+    _extract_metrics,
+    get_machine,
+    make_configs,
+    run_benchmark,
+)
+from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+
+
+class TestBenchRun(ExtTestCase):
+    def test_reg(self):
+        text = ":m,6;"
+        m = _extract_metrics(text)
+        self.assertEqual(m, {"m": 6})
+
+    def test_cmd(self):
+        cmd = _cmd_line("l", m=6)
+        self.assertEqual(cmd[1:], ["-m", "l", "--m", "6"])
+
+    def test_machine(self):
+        ma = get_machine()
+        self.assertIn("machine", ma)
+        self.assertIn("processor", ma)
+        self.assertIn("cpu", ma)
+        self.assertIn("has_cuda", ma)
+        self.assertIn("processor_name", ma)
+
+    def test_run_script(self):
+        script_name = "onnx_diagnostic._bench_test"
+        configs = [dict(m=6)]
+        try:
+            res = run_benchmark(script_name, configs)
+        except BenchmarkError as e:
+            raise unittest.SkipTest(f"Probably no metric collected due to {e}")  # noqa: B904
+        self.assertEqual(len(res), 1)
+        expected = {"metric1": 0.5, "metric2": 5, "metric3": "dummy", "m": 6}
+        got = res[0]
+        for k, v in expected.items():
+            self.assertIn(k, got)
+            self.assertEqual(v, got[k])
+
+    def test_make_configs(self):
+        kwargs = {"single": "1", "multi2": "1,2", "multi3": "A,B,C"}
+        confs = make_configs(kwargs)
+        self.assertEqual(
+            confs,
+            [
+                {"single": "1", "multi2": "1", "multi3": "A"},
+                {"single": "1", "multi2": "1", "multi3": "B"},
+                {"single": "1", "multi2": "1", "multi3": "C"},
+                {"single": "1", "multi2": "2", "multi3": "A"},
+                {"single": "1", "multi2": "2", "multi3": "B"},
+                {"single": "1", "multi2": "2", "multi3": "C"},
+            ],
+        )
+
+    def test_make_configs_last(self):
+        kwargs = {"single": "1", "multi2": "1,2", "multi3": "5, 6"}
+        confs = make_configs(kwargs, last=["multi2"])
+        self.assertEqual(
+            confs,
+            [
+                {"single": "1", "multi3": "5", "multi2": "1"},
+                {"single": "1", "multi3": "5", "multi2": "2"},
+                {"single": "1", "multi3": " 6", "multi2": "1"},
+                {"single": "1", "multi3": " 6", "multi2": "2"},
+            ],
+        )
+
+    def test_make_configs_filter(self):
+        def filter_out(kwargs):
+            if kwargs["multi2"] == "1" and kwargs["multi3"] == "A":
+                return True
+            return False
+
+        kwargs = {"single": "1", "multi2": "1,2", "multi3": "A,B,C"}
+        confs = make_configs(kwargs, filter_function=filter_out)
+        self.assertEqual(confs, [{"single": "1", "multi2": "1", "multi3": "A"}])
+
+    def test_make_configs_drop(self):
+        kwargs = {"single": "1", "multi2": "1,2", "multi3": "A,B,C"}
+        confs = make_configs(kwargs, drop=["multi3"])
+        self.assertEqual(
+            confs,
+            [{"single": "1", "multi2": "1"}, {"single": "1", "multi2": "2"}],
+        )
+
+    def test_make_configs_replace(self):
+        kwargs = {"single": "1", "multi2": "1,2", "multi3": "A,B,C"}
+        confs = make_configs(kwargs, replace={"multi3": "ZZZ"})
+        self.assertEqual(
+            confs,
+            [
+                {"single": "1", "multi2": "1", "multi3": "ZZZ"},
+                {"single": "1", "multi2": "2", "multi3": "ZZZ"},
+            ],
+        )
+
+    def test_max_diff(self):
+        self.assertEqual(
+            max_diff(torch.Tensor([1, 2]), torch.Tensor([1, 2])),
+            {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 2.0, "dnan": 0.0},
+        )
+        self.assertEqual(
+            max_diff(
+                (torch.Tensor([1, 2]),),
+                (torch.Tensor([1, 2])),
+            ),
+            {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 2.0, "dnan": 0.0},
+        )
+        self.assertEqual(
+            max_diff(
+                (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),
+                (torch.Tensor([1, 2]), (torch.Tensor([1, 2]),)),
+            ),
+            {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 4.0, "dnan": 0.0},
+        )
+        self.assertEqual(
+            max_diff(
+                {"a": torch.Tensor([1, 2])},
+                {"a": torch.Tensor([1, 2])},
+            ),
+            {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 2.0, "dnan": 0.0},
+        )
+        self.assertEqual(
+            max_diff(
+                {"a": torch.Tensor([1, 2])},
+                [torch.Tensor([1, 2])],
+            ),
+            {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 2.0, "dnan": 0.0},
+        )
+        self.assertEqual(
+            max_diff(
+                {"a": torch.Tensor([1, float("nan")])},
+                [torch.Tensor([1, 2])],
+            ),
+            {
+                "abs": 9999999998.0,
+                "dnan": 1.0,
+                "n": 2.0,
+                "rel": 0.9999999997999001,
+                "sum": 9999999998.0,
+            },
+        )
+
+    @hide_stdout()
+    def test_max_diff_dynamic_cache(self):
+        t1 = torch.tensor([0, 1], dtype=torch.float32)
+        cache = make_dynamic_cache([(torch.ones((2, 2)), (torch.ones((2, 2)) * 2))])
+        md = max_diff(
+            (t1, cache),
+            (t1, cache.key_cache[0], cache.value_cache[0]),
+            flatten=True,
+            verbose=10,
+        )
+        self.assertEqual(md, {"abs": 0.0, "rel": 0.0, "sum": 0.0, "n": 10.0, "dnan": 0})
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_helpers/test_memory_peak.py

@@ -0,0 +1,76 @@
+import os
+import time
+import unittest
+import numpy as np
+import torch
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    skipif_ci_apple,
+    ignore_warnings,
+    requires_cuda,
+)
+from onnx_diagnostic.helpers.memory_peak import get_memory_rss, start_spying_on
+
+
+class TestMemoryPeak(ExtTestCase):
+    @skipif_ci_apple("stuck")
+    def test_memory(self):
+        mem = get_memory_rss(os.getpid())
+        self.assertIsInstance(mem, int)
+
+    @skipif_ci_apple("stuck")
+    @ignore_warnings(DeprecationWarning)
+    def test_spy_cpu(self):
+        p = start_spying_on(cuda=False)
+        n_elements = 0
+        for _i in range(10):
+            time.sleep(0.005)
+            value = np.empty(2**23, dtype=np.int64)
+            time.sleep(0.005)
+            value += 1
+            time.sleep(0.005)
+            n_elements = max(value.shape[0], n_elements)
+        time.sleep(0.02)
+        pres = p.stop()
+        self.assertGreater(n_elements, 0)
+        self.assertIsInstance(pres, dict)
+        self.assertLessEqual(pres["cpu"].end, pres["cpu"].max_peak)
+        self.assertLessEqual(pres["cpu"].begin, pres["cpu"].max_peak)
+        self.assertGreater(pres["cpu"].begin, 0)
+        # Zero should not happen...
+        self.assertGreaterOrEqual(pres["cpu"].delta_peak, 0)
+        self.assertGreaterOrEqual(pres["cpu"].delta_peak, pres["cpu"].delta_end)
+        self.assertGreaterOrEqual(pres["cpu"].delta_peak, pres["cpu"].delta_avg)
+        self.assertGreaterOrEqual(pres["cpu"].delta_end, 0)
+        self.assertGreaterOrEqual(pres["cpu"].delta_avg, 0)
+        # Too unstable.
+        # self.assertGreater(pres["cpu"].delta_peak, n_elements * 8 * 0.5)
+        self.assertIsInstance(pres["cpu"].to_dict(), dict)
+
+    @skipif_ci_apple("stuck")
+    @requires_cuda()
+    def test_spy_cuda(self):
+        p = start_spying_on(cuda=True)
+        n_elements = 0
+        for _i in range(10):
+            time.sleep(0.005)
+            value = torch.empty(2**23, dtype=torch.int64, device="cuda")
+            value += 1
+            n_elements = max(value.shape[0], n_elements)
+        time.sleep(0.02)
+        pres = p.stop()
+        self.assertIsInstance(pres, dict)
+        self.assertIn("gpus", pres)
+        gpu = pres["gpus"][0]
+        self.assertLessEqual(gpu.end, gpu.max_peak)
+        self.assertLessEqual(gpu.begin, gpu.max_peak)
+        self.assertGreater(gpu.delta_peak, 0)
+        self.assertGreaterOrEqual(gpu.delta_peak, gpu.delta_end)
+        self.assertGreaterOrEqual(gpu.delta_peak, gpu.delta_avg)
+        self.assertGreater(gpu.delta_end, 0)
+        self.assertGreater(gpu.delta_avg, 0)
+        self.assertGreater(gpu.delta_peak, n_elements * 8 * 0.5)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)