misc: Update memory report format

Author: EdCaunt

devito/operator/operator.py

@@ -883,30 +883,38 @@ def estimate_memory(self, human_readable=True, **kwargs):
         args = self._prepare_arguments(estimate_memory=True, **kwargs)
         mem = args.nbytes_consumed
 
+        # Extra information for enhanced operators
+        extras = self._enrich_memreport(args, human_readable=human_readable)
+
         if human_readable:
             headline = f"Memory consumption for operator `{self.name}`:"
             w = len(headline)
             # Columns are width 10
-            fdisk = str(humanbytes(mem[disk_layer])).center(10)
             fhost = str(humanbytes(mem[host_layer])).center(10)
             fdevice = str(humanbytes(mem[device_layer])).center(10)
 
-            # TODO: There is nominally a table generator in pytools which is a dependency
-            # of a dependency and thus in the env anyway
-            info(
+            memreport = (
                 "\n"
                 f"{headline}\n"
-                f"{'┌──────────┬──────────┬──────────┐'.center(w)}\n"
-                f"{'│   Disk   │   Host   │  Device  │'.center(w)}\n"
-                f"{'├──────────┼──────────┼──────────┤'.center(w)}\n"
-                f"{f'│{fdisk}│{fhost}│{fdevice}│'.center(w)}\n"
-                f"{'└──────────┴──────────┴──────────┘'.center(w)}\n"
+                f"{'┌──────────┬──────────┐'.center(w)}\n"
+                f"{'│   Host   │  Device  │'.center(w)}\n"
+                f"{'├──────────┼──────────┤'.center(w)}\n"
+                f"{f'│{fhost}│{fdevice}│'.center(w)}\n"
+                f"{'└──────────┴──────────┘'.center(w)}\n"
             )
 
             # TODO: add hinting if the specified operator won't fit
-
         else:
-            info(f"{self.name} {mem[disk_layer]} {mem[host_layer]} {mem[device_layer]}")
+            memreport = f"{self.name} {mem[host_layer]} {mem[device_layer]}"
+
+        if extras is not None:
+            memreport += extras
+
+        info(memreport)
+
+    def _enrich_memreport(self, args, human_readable=True):
+        # Hook for enriching memory report
+        pass
 
     def apply(self, **kwargs):
         """
@@ -1326,8 +1334,6 @@ def nbytes_avail_mapper(self):
 
         return mapper
 
-    # TODO: This will want some suitable tests in due course
-    # TODO: Might want to also check the spillover onto disk
     @cached_property
     def nbytes_consumed(self):
         """Memory consumed by all objects in the operator"""
@@ -1376,8 +1382,6 @@ def get_nbytes(obj):
                       and not i.is_ArrayBasic and not i.alias]
 
         for i in op_symbols:
-            # FIXME: Probably wrong for streamed functions
-            # TODO: Need a hook for PRO here
             # Will overreport memory usage currently
             try:
                 # TODO: is _obj even needed?
@@ -1465,6 +1469,27 @@ def nbytes_consumed_memmapped(self):
 
         return {disk_layer: 0, host_layer: 0, device_layer: device}
 
+    @cached_property
+    def nbytes_snapshots(self):
+        disk = 0
+
+        # Symbols in the operator which may or may not carry data
+        op_symbols = FindSymbols().visit(self.op)
+
+        # Filter to streamed functions
+        op_symbols = [i for i in op_symbols if i.is_AbstractFunction
+                      and not i.is_ArrayBasic and not i.alias]
+
+        for i in op_symbols:
+            try:
+                disk += i.size_snapshot
+            except AttributeError:
+                pass
+
+        print(disk)
+
+        return {disk_layer: 0, host_layer: 0, device_layer: 0}
+
 
 def parse_kwargs(**kwargs):
     """

tests/test_operator.py

@@ -2073,8 +2073,8 @@ def parse_output(self, output, expected):
         # Check that no allocation occurs as estimate_memory should avoid data touch
         assert "Allocating" not in output.text
 
-        name, disk, host, device = output.records[-1].message.split()
-        extracted = (name, int(disk), int(host), int(device))
+        name, host, device = output.records[-1].message.split()
+        extracted = (name, int(host), int(device))
 
         assert extracted == expected
 
@@ -2093,7 +2093,7 @@ def test_basic_usage(self, caplog, shape, dtype, so):
 
             # Check output of estimate_memory
             host = reduce(mul, f.shape_allocated)*np.dtype(f.dtype).itemsize
-            expected = ("Kernel", 0, host, 0)
+            expected = ("Kernel", host, 0)
             self.parse_output(caplog, expected)
 
     def test_multiple_objects(self, caplog):
@@ -2107,7 +2107,7 @@ def test_multiple_objects(self, caplog):
 
             check = sum(reduce(mul, func.shape_allocated)*np.dtype(func.dtype).itemsize
                         for func in (f, g))
-            expected = ("Kernel", 0, check, 0)
+            expected = ("Kernel", check, 0)
             self.parse_output(caplog, expected)
 
     @pytest.mark.parametrize('time', [True, False])
@@ -2126,7 +2126,7 @@ def test_sparse(self, caplog, time):
 
             check = sum(reduce(mul, func.shape_allocated)*np.dtype(func.dtype).itemsize
                         for func in (f, src, src.coordinates))
-            expected = ("Kernel", 0, check, 0)
+            expected = ("Kernel", check, 0)
             self.parse_output(caplog, expected)
 
     @pytest.mark.parametrize('save', [None, Buffer(3), 10])
@@ -2138,7 +2138,7 @@ def test_timefunction(self, caplog, save):
             op = Operator(Eq(f, 1))
             op.estimate_memory(human_readable=False)
             check = reduce(mul, f.shape_allocated)*np.dtype(f.dtype).itemsize
-            expected = ("Kernel", 0, check, 0)
+            expected = ("Kernel", check, 0)
             self.parse_output(caplog, expected)
 
     def test_mashup(self, caplog):
@@ -2162,7 +2162,7 @@ def test_mashup(self, caplog):
             check = sum(reduce(mul, func.shape_allocated)*np.dtype(func.dtype).itemsize
                         for func in (f, g, src0, src0.coordinates,
                                      src1, src1.coordinates))
-            expected = ("Kernel", 0, check, 0)
+            expected = ("Kernel", check, 0)
             self.parse_output(caplog, expected)
 
     def test_temp_array(self, caplog):
@@ -2194,7 +2194,7 @@ def test_temp_array(self, caplog):
             # Factor in the temp array
             check += reduce(mul, b.shape_allocated)*np.dtype(a.dtype).itemsize
 
-            expected = ("Kernel", 0, check, 0)
+            expected = ("Kernel", check, 0)
             self.parse_output(caplog, expected)
 
     def test_overrides(self, caplog):
@@ -2226,7 +2226,7 @@ def test_overrides(self, caplog):
             check = sum(reduce(mul, func.shape_allocated)*np.dtype(func.dtype).itemsize
                         for func in (f1, tf1, s1, s1.coordinates, st1, st1.coordinates))
 
-            expected = ("Kernel", 0, check, 0)
+            expected = ("Kernel", check, 0)
             self.parse_output(caplog, expected)
 
     def test_device(self, caplog):
@@ -2249,5 +2249,5 @@ def test_device(self, caplog):
             check = reduce(mul, f.shape_allocated)*np.dtype(f.dtype).itemsize
 
             # Matching memory allocated both on host and device for memmap
-            expected = ("Kernel", 0, check, check)
+            expected = ("Kernel", check, check)
             self.parse_output(caplog, expected)