Refactor benchmark_code_compile.py to have job running utility

This provides a common Python interfaces for monitoring resource usage
of subprocesses

Co-authored-by: Schuyler Eldridge <schuyler.eldridge@ibm.com>

2 files changed, 130 insertions, 71 deletions

diff --git a/benchmark/scripts/benchmark_cold_compile.py b/benchmark/scripts/benchmark_cold_compile.py
index aa452d83..c480987c 100755
--- a/benchmark/scripts/benchmark_cold_compile.py
+++ b/benchmark/scripts/benchmark_cold_compile.py
@@ -2,14 +2,14 @@
 # See LICENSE for license details.
 
 import subprocess
-import re
 from statistics import median, stdev
-import sys
 import argparse
 from collections import OrderedDict
 import os
 import numbers
 
+from monitor_job import monitor_job
+
 # Currently hardcoded
 def get_firrtl_repo():
     cmd = ['git', 'rev-parse', '--show-toplevel']
@@ -19,65 +19,13 @@ def get_firrtl_repo():
 
 firrtl_repo = get_firrtl_repo()
 
-platform = ""
-if sys.platform == 'darwin':
-    print("Running on MacOS")
-    platform = 'macos'
-elif sys.platform.startswith("linux"):
-    print("Running on Linux")
-    platform = 'linux'
-else :
-    raise Exception('Unrecognized platform ' + sys.platform)
-
-def time():
-    if platform == 'macos':
-        return ['/usr/bin/time', '-l']
-    if platform == 'linux':
-        return ['/usr/bin/time', '-v']
-
-def extract_max_size(output):
-    regex = ''
-    if platform == 'macos':
-        regex = '(\d+)\s+maximum resident set size'
-    if platform == 'linux':
-        regex = 'Maximum resident set size[^:]*:\s+(\d+)'
-
-    m = re.search(regex, output, re.MULTILINE)
-    if m :
-        return int(m.group(1))
-    else :
-        raise Exception('Max set size not found!')
-
-def extract_run_time(output):
-    regex = ''
-    res = None
-    if platform == 'macos':
-        regex = '(\d+\.\d+)\s+real'
-    if platform == 'linux':
-        regex = 'Elapsed \(wall clock\) time \(h:mm:ss or m:ss\): ([0-9:.]+)'
-    m = re.search(regex, output, re.MULTILINE)
-    if m :
-        text = m.group(1)
-        if platform == 'macos':
-            return float(text)
-        if platform == 'linux':
-            parts = text.split(':')
-            if len(parts) == 3:
-                return float(parts[0]) * 3600 + float(parts[1]) * 60 + float(parts[0])
-            if len(parts) == 2:
-                return float(parts[0]) * 60 + float(parts[1])
-    raise Exception('Runtime not found!')
-
 def run_firrtl(java, jar, design):
     java_cmd = java.split()
-    cmd = time() + java_cmd + ['-cp', jar, 'firrtl.stage.FirrtlMain', '-i', design,'-o','out.v','-X','verilog']
-    result = subprocess.run(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
-    if result.returncode != 0 :
-        print(result.stdout)
-        print(result.stderr)
-        sys.exit(1)
-    size = extract_max_size(result.stderr.decode('utf-8'))
-    runtime = extract_run_time(result.stderr.decode('utf-8'))
+    cmd = java_cmd + ['-cp', jar, 'firrtl.stage.FirrtlMain', '-i', design,'-o','out.v','-X','verilog']
+    print(' '.join(cmd))
+    resource_use = monitor_job(cmd)
+    size = resource_use.maxrss // 1024 # KiB -> MiB
+    runtime = resource_use.wall_clock_time
     return (size, runtime)
 
 def parseargs():
@@ -138,15 +86,6 @@ def check_designs(designs):
     for design in designs:
         assert os.path.exists(design), '{} must be an existing file!'.format(design)
 
-# /usr/bin/time -v on Linux returns size in kbytes
-# /usr/bin/time -l on MacOS returns size in Bytes
-def norm_max_set_sizes(sizes):
-    div = None
-    if platform == 'linux':
-        d = 1000.0
-    if platform == 'macos':
-        d = 1000000.0
-    return [s / d for s in sizes]
 
 def main():
     args = parseargs()
@@ -156,7 +95,7 @@ def main():
     jars = build_firrtl_jars(hashes)
     jvms = args.jvms
     N = args.iterations
-    info = [['java', 'revision', 'design', 'max heap', 'SD', 'runtime', 'SD']]
+    info = [['java', 'revision', 'design', 'max heap (MiB)', 'SD', 'runtime (s)', 'SD']]
     for java in jvms:
         print("Running with '{}'".format(java))
         for hashcode, jar in jars.items():
@@ -166,8 +105,7 @@ def main():
             for design in designs:
                 print('Running {}...'.format(design))
                 (sizes, runtimes) = zip(*[run_firrtl(java, jar, design) for i in range(N)])
-                norm_sizes = norm_max_set_sizes(sizes)
-                info.append([java_title, revision, design, median(norm_sizes), stdev(norm_sizes), median(runtimes), stdev(runtimes)])
+                info.append([java_title, revision, design, median(sizes), stdev(sizes), median(runtimes), stdev(runtimes)])
                 java_title = ''
                 revision = ''
 
diff --git a/benchmark/scripts/monitor_job.py b/benchmark/scripts/monitor_job.py
new file mode 100644
index 00000000..b3e79ead
--- /dev/null
+++ b/benchmark/scripts/monitor_job.py
@@ -0,0 +1,121 @@
+# See LICENSE for license details.
+
+"""
+Utilities for running a subprocess and collecting runtime and memory use information.
+"""
+
+from typing import NamedTuple, List
+from enum import Enum
+import sys
+import subprocess
+import re
+
+# All times in seconds, sizes in kibibytes (KiB)
+JobResourceUse = NamedTuple('JobResourceUse', [('user_time', float),
+                                               ('system_time', float),
+                                               ('wall_clock_time', float),
+                                               ('maxrss', int)])
+
+
+class JobFailedError(Exception):
+    pass
+
+
+def monitor_job(args: List[str]) -> JobResourceUse:
+    """Run a job with resource monitoring, returns resource usage"""
+    platform = get_platform()
+    cmd = time(platform) + args
+    result = subprocess.run(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
+    if result.returncode != 0 :
+        msg = "[stdout]\n{}\n[stderr]\n{}".format(result.stdout, result.stderr)
+        raise JobFailedError(msg)
+    stderr = result.stderr.decode('utf-8')
+
+    user_time =       extract_user_time(platform, stderr)
+    system_time =     extract_system_time(platform, stderr)
+    wall_clock_time = extract_wall_clock_time(platform, stderr)
+    maxrss =          extract_maxrss(platform, stderr)
+    return JobResourceUse(user_time, system_time, wall_clock_time, maxrss)
+
+
+Platform = Enum('Platform', 'linux macos')
+
+
+def get_platform() -> Platform:
+    p = sys.platform
+    if p == 'darwin':
+        return Platform.macos
+    elif p.startswith('linux'):
+        return Platform.linux
+    else:
+        raise Exception('Unsupported platform: ' + p)
+
+
+def time(platform):
+    return { Platform.macos: ['/usr/bin/time', '-l'],
+             Platform.linux: ['/usr/bin/time', '-v'] }[platform]
+
+
+def extract_maxrss(platform, output):
+    """Returns maxrss in kbytes"""
+    regex = {
+        Platform.macos: '(\d+)\s+maximum resident set size',
+        Platform.linux: 'Maximum resident set size[^:]*:\s+(\d+)'
+    }[platform]
+
+    m = re.search(regex, output, re.MULTILINE)
+    if m :
+        return {
+            # /usr/bin/time -l on MacOS returns size in Bytes
+            Platform.macos: int(m.group(1)) // 1024,
+            # /usr/bin/time -v on Linux returns size in kibibytes
+            # https://stackoverflow.com/questions/61392725/kilobytes-or-kibibytes-in-gnu-time
+            Platform.linux: int(m.group(1))
+        }[platform]
+    else :
+        raise Exception('Max set size not found!')
+
+
+def extract_user_time(platform, output):
+    res = None
+    regex = {
+        Platform.macos: '(\d+\.\d+)\s+user',
+        Platform.linux: 'User time \(seconds\): (\d+\.\d+)'
+    }[platform]
+    m = re.search(regex, output, re.MULTILINE)
+    if m :
+        return float(m.group(1))
+    raise Exception('User time not found!')
+
+
+def extract_system_time(platform, output):
+    res = None
+    regex = {
+        Platform.macos: '(\d+\.\d+)\s+sys',
+        Platform.linux: 'System time \(seconds\): (\d+\.\d+)'
+    }[platform]
+    m = re.search(regex, output, re.MULTILINE)
+    if m :
+        return float(m.group(1))
+    raise Exception('System time not found!')
+
+
+def extract_wall_clock_time(platform, output):
+    res = None
+    regex = {
+        Platform.macos: '(\d+\.\d+)\s+real',
+        Platform.linux: 'Elapsed \(wall clock\) time \(h:mm:ss or m:ss\): ([0-9:.]+)'
+    }[platform]
+    m = re.search(regex, output, re.MULTILINE)
+    if m :
+        text = m.group(1)
+        if platform == Platform.macos:
+            return float(text)
+        if platform == Platform.linux:
+            parts = text.split(':')
+            if len(parts) == 3:
+                return float(parts[0]) * 3600 + float(parts[1]) * 60 + float(parts[0])
+            if len(parts) == 2:
+                return float(parts[0]) * 60 + float(parts[1])
+    raise Exception('Wall clock time not found!')
+