misc: Update memory report format

Author: EdCaunt

devito/operator/operator.py

@@ -630,7 +630,6 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
             if p.name in args:
                 # E.g., SubFunctions
                 continue
-            # print(p._arg_values(**kwargs))  # Trigger first-touch
             for k, v in p._arg_values(estimate_memory=estimate_memory, **kwargs).items():
                 if k not in args:
                     args[k] = v
@@ -659,7 +658,6 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
         # the subsequent phases of the arguments processing
         args = kwargs['args'] = ArgumentsMap(args, grid, self)
 
-        # FIXME: Will want to remove this if not using prepare_args to estimate memory
         if estimate_memory:
             # No need to do anything more if only checking the memory
             return args
@@ -888,30 +886,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):
         """
@@ -1338,7 +1344,6 @@ def nbytes_avail_mapper(self):
         """
         mapper = {}
 
-        # TODO: This doesn't account for the size of the snapshots?
         # The amount of space available on the disk
         usage = shutil.disk_usage(gettempdir())
         mapper[disk_layer] = usage.free
@@ -1356,17 +1361,12 @@ def nbytes_avail_mapper(self):
         mapper[host_layer] = int(ANYCPU.memavail() / nproc)
 
         for layer in (host_layer, device_layer):
-            try:
-                mapper[layer] -= self.nbytes_consumed_operator[layer]
-            except KeyError:
-                continue
+            mapper[layer] -= self.nbytes_consumed_operator.get(layer, 0)
 
         mapper = {k: int(v) for k, v in mapper.items()}
 
         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"""
@@ -1415,17 +1415,13 @@ 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?
                 v = get_nbytes(self[i.name]._obj)
             except AttributeError:
                 v = get_nbytes(self.get(i.name, i))
 
-            print(i, humanbytes(v))
-
             if i._mem_host or i._mem_mapped:
                 # No need to add to device , as it will be counted
                 # by nbytes_consumed_memmapped
@@ -1462,8 +1458,6 @@ def nbytes_consumed_array(self):
                 # E.g. the Arrays used to store the MPI halo exchanges
                 continue
 
-            print(i, humanbytes(v))
-
             if i._mem_host:
                 host += v
             elif i._mem_local:
@@ -1504,6 +1498,25 @@ def nbytes_consumed_memmapped(self):
 
         return {disk_layer: 0, host_layer: 0, device_layer: device}
 
+    @cached_property
+    def nbytes_snapshots(self):
+
+        # 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]
+
+        disk = 0
+        for i in op_symbols:
+            try:
+                disk += i.size_snapshot*i._time_size_ideal*np.dtype(i.dtype).itemsize
+            except AttributeError:
+                pass
+
+        return {disk_layer: disk, host_layer: 0, device_layer: 0}
+
 
 def parse_kwargs(**kwargs):
     """

devito/types/sparse.py

@@ -652,7 +652,7 @@ def _arg_defaults(self, alias=None, estimate_memory=False):
             if f is not None:
                 mapper[getattr(self, i)] = f
 
-        if estimate_memory:  # kwargs.get("estimate_memory", False):
+        if estimate_memory:
             # Avoid touching the data in any capacity, and simply return
             # the symbolic objects if merely estimating memory consumption.
             return ReducerMap({v.name: k for k, v in mapper.items()})

tests/test_operator.py

@@ -2073,8 +2073,9 @@ 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))
+        parsed = output.records[-1].message.split()
+        name, host, device = parsed[:3]
+        extracted = (name, int(host), int(device))
 
         assert extracted == expected
 
@@ -2093,7 +2094,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 +2108,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 +2127,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 +2139,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 +2163,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 +2195,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 +2227,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 +2250,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)