dsl: Add method to dump memory estimate to JSON

Author: EdCaunt

devito/tools/data_structures.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict, deque
 from collections.abc import Callable, Iterable, MutableSet, Mapping, Set
 from functools import reduce, cached_property
+import json
 
 import numpy as np
 from multidict import MultiDict
@@ -672,6 +673,11 @@ class MemoryEstimate(frozendict):
     human_readable: frozendict
         The mapper, albeit with human-readable memory usage (MB, GB, etc)
         rather than raw bytes.
+
+    Methods
+    -------
+    to_json(path)
+        Write the memory estimate to a JSON for scheduler ingestion.
     """
 
     def __init__(self, *args, **kwargs):
@@ -690,6 +696,21 @@ def human_readable(self):
     def __repr__(self):
         return f'{self.__class__.__name__}({self.name}): {self.human_readable._dict}'
 
+    def to_json(self, path):
+        """
+        Write the MemoryEstimate to JSON for ingestion by a scheduler.
+
+        Arguments
+        ---------
+        path: str
+            The path to which the memory estimate should be written.
+        """
+        summary = {'name': self.name, **self._dict}
+        json_object = json.dumps(summary, indent=4)
+
+        with open(path, "w") as outfile:
+            outfile.write(json_object)
+
 
 class UnboundTuple(tuple):
     """

tests/test_operator.py

@@ -1,7 +1,9 @@
+import os
 from itertools import permutations
 from functools import reduce
 from operator import mul
 import logging
+import json
 
 import numpy as np
 import sympy
@@ -2268,3 +2270,21 @@ def test_device(self, caplog):
 
             # Matching memory allocated both on host and device for memmap
             self.parse_output(summary, check)
+
+    def test_to_json(self):
+        grid = Grid(shape=(101, 101))
+        f = Function(name='f', grid=grid, space_order=2)
+        op = Operator(Eq(f, 1))
+        summary = op.estimate_memory()
+
+        summary.to_json("memory_estimate_output.json")
+
+        with open("memory_estimate_output.json", "r") as infile:
+            json_object = json.load(infile)
+
+            assert json_object['name'] == summary.name
+            assert json_object['host'] == summary['host']
+            assert json_object['device'] == summary['device']
+
+        # Clean up
+        os.remove("memory_estimate_output.json")